WIP: compress benchmark

build.gradle.kts

@@ -51,6 +51,7 @@ allprojects {
   version = project.property("VERSION_NAME") as String
 
   repositories {
+    mavenLocal()
     mavenCentral()
     google()
   }

gradle/libs.versions.toml

@@ -5,7 +5,7 @@ ktlint = "0.48.2"
 
 [libraries]
 android-desugar-jdk-libs = { module = "com.android.tools:desugar_jdk_libs", version = "2.1.5" }
-android-gradle-plugin = { module = "com.android.tools.build:gradle", version = "8.11.0" }
+android-gradle-plugin = { module = "com.android.tools.build:gradle", version = "8.10.0" }
 androidx-test-ext-junit = { module = "androidx.test.ext:junit", version = "1.2.1" }
 androidx-test-runner = { module = "androidx.test:runner", version = "1.5.2" }
 binaryCompatibilityValidator = { module = "org.jetbrains.kotlinx.binary-compatibility-validator:org.jetbrains.kotlinx.binary-compatibility-validator.gradle.plugin", version = "0.18.0" }

okio/jvm/jmh/build.gradle.kts

@@ -4,9 +4,12 @@ plugins {
 }
 
 jmh {
+  includes.add("ZstdImplementationBenchmark")
 }
 
 dependencies {
   api(projects.okio)
   api(libs.jmh.core)
+  api("com.squareup.okio-zstd:okio-zstd:1.0.0-SNAPSHOT")
+  api("com.github.luben:zstd-jni:1.5.7-4")
 }

okio/jvm/jmh/src/jmh/java/com/squareup/okio/benchmarks/CompressBenchmark.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2025 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.okio.benchmarks;
+
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+import okio.Buffer;
+import okio.BufferedSink;
+import okio.BufferedSource;
+import okio.Okio;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+@Fork(1)
+@Warmup(iterations = 1, time = 2)
+@Measurement(iterations = 3, time = 2)
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+public class CompressBenchmark {
+  private SampleData sampleData;
+
+  @Param({"deflate", "gzip", "zstd", "none"})
+  String algorithm;
+
+  @Param({"8388608"}) // 8 MiB.
+  int sampleDataSize;
+
+  @Setup
+  public void setup() throws IOException {
+    sampleData = SampleData.create(algorithm, sampleDataSize);
+  }
+
+  @Benchmark
+  public void compress() throws IOException {
+    try (BufferedSink sink = Okio.buffer(sampleData.compress(algorithm, Okio.blackhole()))) {
+      sink.write(sampleData.uncompressedData);
+    }
+  }
+
+  @Benchmark
+  public void decompress() throws IOException {
+    Buffer compressedBuffer = new Buffer();
+    compressedBuffer.write(sampleData.compressedData);
+    try (BufferedSource source = Okio.buffer(sampleData.decompress(algorithm, compressedBuffer))) {
+      source.readAll(Okio.blackhole());
+    }
+  }
+}

okio/jvm/jmh/src/jmh/java/com/squareup/okio/benchmarks/SampleData.java

@@ -0,0 +1,101 @@
+package com.squareup.okio.benchmarks;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.zip.Deflater;
+import java.util.zip.Inflater;
+import okio.Buffer;
+import okio.BufferedSink;
+import okio.BufferedSource;
+import okio.DeflaterSink;
+import okio.FileSystem;
+import okio.GzipSink;
+import okio.GzipSource;
+import okio.InflaterSource;
+import okio.Okio;
+import okio.Path;
+import okio.Sink;
+import okio.Source;
+import okio.zstd.Zstd;
+
+public class SampleData {
+  private static final Path root = Path.get("/Volumes/Development/cash-android", false);
+  private static final FileSystem fileSystem = FileSystem.SYSTEM;
+
+  public final byte[] uncompressedData;
+  public final byte[] compressedData;
+
+  public SampleData(byte[] uncompressedData, byte[] compressedData) {
+    this.uncompressedData = uncompressedData;
+    this.compressedData = compressedData;
+  }
+
+  public static SampleData create(String algorithm, int size) throws IOException {
+    byte[] uncompressedData = generateSampleData(size);
+
+    Buffer compressedBuffer = new Buffer();
+    try (BufferedSink sink = Okio.buffer(compress(algorithm, compressedBuffer))) {
+      sink.write(uncompressedData);
+    }
+    byte[] compressedData = compressedBuffer.readByteArray();
+
+    // Confirm the compression is round-trip.
+    Buffer validateSource = new Buffer();
+    validateSource.write(compressedData);
+    byte[] decompressedData;
+    try (BufferedSource source = Okio.buffer(decompress(algorithm, validateSource))) {
+      decompressedData = source.readByteArray();
+    }
+
+    if (!Arrays.equals(uncompressedData, decompressedData)) {
+      throw new IllegalStateException("failed to round trip " + algorithm);
+    }
+
+    return new SampleData(uncompressedData, compressedData);
+  }
+
+  private static byte[] generateSampleData(int size) throws IOException {
+    Buffer sampleDataBuffer = new Buffer();
+    Iterator<Path> pathIterator = fileSystem.listRecursively(root).iterator();
+    while (sampleDataBuffer.size() < size) {
+      Path path = pathIterator.next();
+      if (path.name().endsWith(".kt")) {
+        Source source = fileSystem.source(path);
+        try {
+          sampleDataBuffer.writeAll(source);
+        } finally {
+          source.close();
+        }
+      }
+    }
+    return sampleDataBuffer.readByteArray(size);
+  }
+  static Sink compress(String algorithm, Sink delegate) {
+    if (algorithm.equals("deflate")) {
+      return new DeflaterSink(delegate, new Deflater());
+    } else if (algorithm.equals("gzip")) {
+      return new GzipSink(delegate);
+    } else if (algorithm.equals("zstd")) {
+      return Zstd.zstdCompress(delegate);
+    } else if (algorithm.equals("none")) {
+      return delegate;
+    } else {
+      throw new IllegalArgumentException("unexpected algorithm: " + algorithm);
+    }
+  }
+
+  static Source decompress(String algorithm, Source delegate) {
+    if (algorithm.equals("deflate")) {
+      return new InflaterSource(delegate, new Inflater());
+    } else if (algorithm.equals("gzip")) {
+      return new GzipSource(delegate);
+    } else if (algorithm.equals("zstd")) {
+      return Zstd.zstdDecompress(delegate);
+    } else if (algorithm.equals("none")) {
+      return delegate;
+    } else {
+      throw new IllegalArgumentException("unexpected algorithm: " + algorithm);
+    }
+  }
+}

okio/jvm/jmh/src/jmh/java/com/squareup/okio/benchmarks/ZstdImplementationBenchmark.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2025 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.okio.benchmarks;
+
+import com.github.luben.zstd.ZstdInputStream;
+import com.github.luben.zstd.ZstdOutputStream;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.concurrent.TimeUnit;
+import okio.Buffer;
+import okio.BufferedSink;
+import okio.BufferedSource;
+import okio.Okio;
+import org.jetbrains.annotations.NotNull;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+/** Confirm Okio Zstd has performance consistent with Zstd-jni. */
+@Fork(1)
+@Warmup(iterations = 1, time = 2)
+@Measurement(iterations = 3, time = 2)
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+public class ZstdImplementationBenchmark {
+  private SampleData sampleData;
+
+  @Param({"8388608"}) // 8 MiB.
+  int sampleDataSize;
+
+  @Setup
+  public void setup() throws IOException {
+    sampleData = SampleData.create("zstd", sampleDataSize);
+  }
+
+  @Benchmark
+  public void okioCompress() throws IOException {
+    try (BufferedSink sink = Okio.buffer(sampleData.compress("zstd", Okio.blackhole()))) {
+      sink.write(sampleData.uncompressedData);
+    }
+  }
+
+  @Benchmark
+  public void zstdJniCompress() throws IOException {
+    try (OutputStream out = new ZstdOutputStream(new BlackholeOutputStream())) {
+      out.write(sampleData.uncompressedData);
+    }
+  }
+
+  @Benchmark
+  public void okioDecompress() throws IOException {
+    Buffer compressedBuffer = new Buffer();
+    compressedBuffer.write(sampleData.compressedData);
+    try (BufferedSource source = Okio.buffer(sampleData.decompress("zstd", compressedBuffer))) {
+      source.readAll(Okio.blackhole());
+    }
+  }
+
+  @Benchmark
+  public void zstdJniDecompress() throws IOException {
+    byte[] blackhole = new byte[8192];
+    InputStream compressedInputStream = new ByteArrayInputStream(sampleData.compressedData);
+    try (InputStream in = new ZstdInputStream(compressedInputStream)) {
+      while (true) {
+        if (in.read(blackhole) == -1) break;
+      }
+    }
+  }
+
+  static class BlackholeOutputStream extends OutputStream {
+    @Override
+    public void write(@NotNull byte[] b, int off, int len) throws IOException {
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+    }
+
+    @Override
+    public void write(@NotNull byte[] b) throws IOException {
+    }
+  }
+}