src/test/java/com/google/devtools/build/lib/remote/ChunkedTransferBenchmark.java - bazel - Git at Google


            
            // Copyright 2026 The Bazel Authors. All rights reserved.

            
            //

            
            // 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.google.devtools.build.lib.remote;


            import static org.mockito.ArgumentMatchers.any;

            
            import static org.mockito.ArgumentMatchers.anyBoolean;

            
            import static org.mockito.Mockito.mock;

            
            import static org.mockito.Mockito.when;


            import build.bazel.remote.execution.v2.Digest;

            
            import build.bazel.remote.execution.v2.RequestMetadata;

            
            import build.bazel.remote.execution.v2.SplitBlobResponse;

            
            import com.google.common.collect.ImmutableSet;

            
            import com.google.common.util.concurrent.Futures;

            
            import com.google.common.util.concurrent.ListenableFuture;

            
            import com.google.common.util.concurrent.SettableFuture;

            
            import com.google.devtools.build.lib.clock.JavaClock;

            
            import com.google.devtools.build.lib.remote.chunking.FastCdcChunker;

            
            import com.google.devtools.build.lib.remote.chunking.FastCdcChunkingConfig;

            
            import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;

            
            import com.google.devtools.build.lib.remote.common.RemoteCacheClient.Blob;

            
            import com.google.devtools.build.lib.remote.util.DigestUtil;

            
            import com.google.devtools.build.lib.vfs.DigestHashFunction;

            
            import com.google.devtools.build.lib.vfs.FileSystem;

            
            import com.google.devtools.build.lib.vfs.Path;

            
            import com.google.devtools.build.lib.vfs.SyscallCache;

            
            import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;

            
            import java.io.OutputStream;

            
            import java.util.ArrayList;

            
            import java.util.HashMap;

            
            import java.util.List;

            
            import java.util.Map;

            
            import java.util.Random;

            
            import java.util.concurrent.Executors;

            
            import java.util.concurrent.ScheduledExecutorService;

            
            import java.util.concurrent.TimeUnit;

            
            import org.openjdk.jmh.annotations.Benchmark;

            
            import org.openjdk.jmh.annotations.BenchmarkMode;

            
            import org.openjdk.jmh.annotations.Fork;

            
            import org.openjdk.jmh.annotations.Level;

            
            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.TearDown;

            
            import org.openjdk.jmh.annotations.Warmup;


            /** Benchmark for chunk download/upload with per-chunk latency jitter. */

            
            @BenchmarkMode(Mode.AverageTime)

            
            @OutputTimeUnit(TimeUnit.MILLISECONDS)

            
            @Warmup(iterations = 1, time = 1, timeUnit = TimeUnit.SECONDS)

            
            @Measurement(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS)

            
            @Fork(1)

            
            public class ChunkedTransferBenchmark {

            
              private static final DigestUtil DIGEST_UTIL =

            
                  new DigestUtil(SyscallCache.NO_CACHE, DigestHashFunction.SHA256);

            
              private static final RemoteActionExecutionContext CONTEXT =

            
                  RemoteActionExecutionContext.create(RequestMetadata.getDefaultInstance());


              @Benchmark

            
              public void downloadChunked(DownloadState state) throws Exception {

            
                state.downloader.downloadChunked(CONTEXT, state.blobDigest, OutputStream.nullOutputStream());

            
              }


              @Benchmark

            
              public void uploadChunked(UploadState state) throws Exception {

            
                state.uploader.uploadChunked(CONTEXT, state.blobDigest, state.file, /* force= */ false);

            
              }


              @State(Scope.Thread)

            
              public static class DownloadState {

            
                @Param({"1", "2", "4", "8"})

            
                public int schedulerThreads;


                @Param({"32"})

            
                public int chunkCount;


                @Param({"1024"})

            
                public int chunkSizeBytes;


                @Param({"25"})

            
                public int delayMillis;


                @Param({"10"})

            
                public int jitterMillis;


                private ScheduledExecutorService scheduler;

            
                private ChunkedBlobDownloader downloader;

            
                private Digest blobDigest;

            
                private Random latencyJitter;


                @Setup(Level.Trial)

            
                public void setup() throws Exception {

            
                  scheduler = Executors.newScheduledThreadPool(schedulerThreads);

            
                  latencyJitter = new Random(12345L);


                  GrpcCacheClient grpcCacheClient = mock(GrpcCacheClient.class);

            
                  CombinedCache combinedCache = mock(CombinedCache.class);


                  List<Digest> chunkDigests = new ArrayList<>(chunkCount);

            
                  Map<Digest, byte[]> chunkDataByDigest = new HashMap<>(chunkCount);

            
                  long totalBytes = 0;

            
                  for (int i = 0; i < chunkCount; i++) {

            
                    byte[] chunkData = new byte[chunkSizeBytes];

            
                    new Random(i).nextBytes(chunkData);

            
                    Digest chunkDigest = DIGEST_UTIL.compute(chunkData);

            
                    chunkDigests.add(chunkDigest);

            
                    chunkDataByDigest.put(chunkDigest, chunkData);

            
                    totalBytes += chunkData.length;

            
                  }


                  when(combinedCache.downloadBlob(any(), any(Digest.class)))

            
                      .thenAnswer(

            
                          invocation ->

            
                              delayedFuture(

            
                                  chunkDataByDigest.get(invocation.getArgument(1)),

            
                                  delayMillis,

            
                                  jitterMillis,

            
                                  latencyJitter,

            
                                  scheduler));


                  blobDigest =

            
                      Digest.newBuilder()

            
                          .setHash("chunked-transfer-benchmark-download-" + chunkCount + "-" + chunkSizeBytes)

            
                          .setSizeBytes(totalBytes)

            
                          .build();


                  SplitBlobResponse splitBlobResponse =

            
                      SplitBlobResponse.newBuilder().addAllChunkDigests(chunkDigests).build();

            
                  when(grpcCacheClient.splitBlob(any(), any(Digest.class), any()))

            
                      .thenReturn(Futures.immediateFuture(splitBlobResponse));


                  FastCdcChunkingConfig chunkingConfig = new FastCdcChunkingConfig(chunkSizeBytes, 2, 0);

            
                  downloader =

            
                      new ChunkedBlobDownloader(

            
                          grpcCacheClient, combinedCache, chunkingConfig, DIGEST_UTIL, new ChunkLocationMap());

            
                }


                @TearDown(Level.Trial)

            
                public void tearDown() {

            
                  scheduler.shutdownNow();

            
                }

            
              }


              @State(Scope.Thread)

            
              public static class UploadState {

            
                @Param({"1", "2", "4", "8"})

            
                public int schedulerThreads;


                @Param({"32768"})

            
                public int fileSizeBytes;


                @Param({"1024"})

            
                public int avgChunkSizeBytes;


                @Param({"25"})

            
                public int delayMillis;


                @Param({"10"})

            
                public int jitterMillis;


                private ScheduledExecutorService scheduler;

            
                private ChunkedBlobUploader uploader;

            
                private Path file;

            
                private Digest blobDigest;

            
                private Random latencyJitter;


                @Setup(Level.Trial)

            
                public void setup() throws Exception {

            
                  scheduler = Executors.newScheduledThreadPool(schedulerThreads);

            
                  latencyJitter = new Random(54321L);


                  GrpcCacheClient grpcCacheClient = mock(GrpcCacheClient.class);

            
                  CombinedCache combinedCache = mock(CombinedCache.class);


                  byte[] data = new byte[fileSizeBytes];

            
                  new Random(42).nextBytes(data);

            
                  blobDigest = DIGEST_UTIL.compute(data);


                  FileSystem fs = new InMemoryFileSystem(new JavaClock(), DigestHashFunction.SHA256);

            
                  file = fs.getPath("/bench/blob.bin");

            
                  file.getParentDirectory().createDirectoryAndParents();

            
                  try (var out = file.getOutputStream()) {

            
                    out.write(data);

            
                  }


                  FastCdcChunkingConfig chunkingConfig = new FastCdcChunkingConfig(avgChunkSizeBytes, 2, 0);

            
                  uploader =

            
                      new ChunkedBlobUploader(grpcCacheClient, combinedCache, chunkingConfig, DIGEST_UTIL);


                  List<Digest> chunkDigests;

            
                  try (var input = file.getInputStream()) {

            
                    chunkDigests = new FastCdcChunker(chunkingConfig, DIGEST_UTIL).chunkToDigests(input);

            
                  }


                  when(grpcCacheClient.findMissingDigests(any(), any()))

            
                      .thenReturn(Futures.immediateFuture(ImmutableSet.copyOf(chunkDigests)));

            
                  when(grpcCacheClient.spliceBlob(any(), any(Digest.class), any(), any()))

            
                      .thenReturn(Futures.immediateVoidFuture());

            
                  when(combinedCache.uploadBlob(any(), any(Digest.class), any(Blob.class), anyBoolean()))

            
                      .thenAnswer(

            
                          invocation ->

            
                              delayedFuture(null, delayMillis, jitterMillis, latencyJitter, scheduler));

            
                }


                @TearDown(Level.Trial)

            
                public void tearDown() {

            
                  scheduler.shutdownNow();

            
                }

            
              }


              private static <T> ListenableFuture<T> delayedFuture(

            
                  T value,

            
                  int delayMillis,

            
                  int jitterMillis,

            
                  Random latencyJitter,

            
                  ScheduledExecutorService scheduler) {

            
                SettableFuture<T> future = SettableFuture.create();

            
                var unused =

            
                    scheduler.schedule(

            
                        () -> future.set(value),

            
                        jitteredDelayMillis(delayMillis, jitterMillis, latencyJitter),

            
                        TimeUnit.MILLISECONDS);

            
                return future;

            
              }


              private static int jitteredDelayMillis(int delayMillis, int jitterMillis, Random latencyJitter) {

            
                if (jitterMillis == 0) {

            
                  return delayMillis;

            
                }

            
                return Math.max(0, delayMillis + latencyJitter.nextInt((jitterMillis * 2) + 1) - jitterMillis);

            
              }

            
            }
	// Copyright 2026 The Bazel Authors. All rights reserved.
	//
	// 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.google.devtools.build.lib.remote;

	import static org.mockito.ArgumentMatchers.any;
	import static org.mockito.ArgumentMatchers.anyBoolean;
	import static org.mockito.Mockito.mock;
	import static org.mockito.Mockito.when;

	import build.bazel.remote.execution.v2.Digest;
	import build.bazel.remote.execution.v2.RequestMetadata;
	import build.bazel.remote.execution.v2.SplitBlobResponse;
	import com.google.common.collect.ImmutableSet;
	import com.google.common.util.concurrent.Futures;
	import com.google.common.util.concurrent.ListenableFuture;
	import com.google.common.util.concurrent.SettableFuture;
	import com.google.devtools.build.lib.clock.JavaClock;
	import com.google.devtools.build.lib.remote.chunking.FastCdcChunker;
	import com.google.devtools.build.lib.remote.chunking.FastCdcChunkingConfig;
	import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
	import com.google.devtools.build.lib.remote.common.RemoteCacheClient.Blob;
	import com.google.devtools.build.lib.remote.util.DigestUtil;
	import com.google.devtools.build.lib.vfs.DigestHashFunction;
	import com.google.devtools.build.lib.vfs.FileSystem;
	import com.google.devtools.build.lib.vfs.Path;
	import com.google.devtools.build.lib.vfs.SyscallCache;
	import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
	import java.io.OutputStream;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Random;
	import java.util.concurrent.Executors;
	import java.util.concurrent.ScheduledExecutorService;
	import java.util.concurrent.TimeUnit;
	import org.openjdk.jmh.annotations.Benchmark;
	import org.openjdk.jmh.annotations.BenchmarkMode;
	import org.openjdk.jmh.annotations.Fork;
	import org.openjdk.jmh.annotations.Level;
	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.TearDown;
	import org.openjdk.jmh.annotations.Warmup;

	/** Benchmark for chunk download/upload with per-chunk latency jitter. */
	@BenchmarkMode(Mode.AverageTime)
	@OutputTimeUnit(TimeUnit.MILLISECONDS)
	@Warmup(iterations = 1, time = 1, timeUnit = TimeUnit.SECONDS)
	@Measurement(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS)
	@Fork(1)
	public class ChunkedTransferBenchmark {
	private static final DigestUtil DIGEST_UTIL =
	new DigestUtil(SyscallCache.NO_CACHE, DigestHashFunction.SHA256);
	private static final RemoteActionExecutionContext CONTEXT =
	RemoteActionExecutionContext.create(RequestMetadata.getDefaultInstance());

	@Benchmark
	public void downloadChunked(DownloadState state) throws Exception {
	state.downloader.downloadChunked(CONTEXT, state.blobDigest, OutputStream.nullOutputStream());
	}

	@Benchmark
	public void uploadChunked(UploadState state) throws Exception {
	state.uploader.uploadChunked(CONTEXT, state.blobDigest, state.file, /* force= */ false);
	}

	@State(Scope.Thread)
	public static class DownloadState {
	@Param({"1", "2", "4", "8"})
	public int schedulerThreads;

	@Param({"32"})
	public int chunkCount;

	@Param({"1024"})
	public int chunkSizeBytes;

	@Param({"25"})
	public int delayMillis;

	@Param({"10"})
	public int jitterMillis;

	private ScheduledExecutorService scheduler;
	private ChunkedBlobDownloader downloader;
	private Digest blobDigest;
	private Random latencyJitter;

	@Setup(Level.Trial)
	public void setup() throws Exception {
	scheduler = Executors.newScheduledThreadPool(schedulerThreads);
	latencyJitter = new Random(12345L);

	GrpcCacheClient grpcCacheClient = mock(GrpcCacheClient.class);
	CombinedCache combinedCache = mock(CombinedCache.class);

	List<Digest> chunkDigests = new ArrayList<>(chunkCount);
	Map<Digest, byte[]> chunkDataByDigest = new HashMap<>(chunkCount);
	long totalBytes = 0;
	for (int i = 0; i < chunkCount; i++) {
	byte[] chunkData = new byte[chunkSizeBytes];
	new Random(i).nextBytes(chunkData);
	Digest chunkDigest = DIGEST_UTIL.compute(chunkData);
	chunkDigests.add(chunkDigest);
	chunkDataByDigest.put(chunkDigest, chunkData);
	totalBytes += chunkData.length;
	}

	when(combinedCache.downloadBlob(any(), any(Digest.class)))
	.thenAnswer(
	invocation ->
	delayedFuture(
	chunkDataByDigest.get(invocation.getArgument(1)),
	delayMillis,
	jitterMillis,
	latencyJitter,
	scheduler));

	blobDigest =
	Digest.newBuilder()
	.setHash("chunked-transfer-benchmark-download-" + chunkCount + "-" + chunkSizeBytes)
	.setSizeBytes(totalBytes)
	.build();

	SplitBlobResponse splitBlobResponse =
	SplitBlobResponse.newBuilder().addAllChunkDigests(chunkDigests).build();
	when(grpcCacheClient.splitBlob(any(), any(Digest.class), any()))
	.thenReturn(Futures.immediateFuture(splitBlobResponse));

	FastCdcChunkingConfig chunkingConfig = new FastCdcChunkingConfig(chunkSizeBytes, 2, 0);
	downloader =
	new ChunkedBlobDownloader(
	grpcCacheClient, combinedCache, chunkingConfig, DIGEST_UTIL, new ChunkLocationMap());
	}

	@TearDown(Level.Trial)
	public void tearDown() {
	scheduler.shutdownNow();
	}
	}

	@State(Scope.Thread)
	public static class UploadState {
	@Param({"1", "2", "4", "8"})
	public int schedulerThreads;

	@Param({"32768"})
	public int fileSizeBytes;

	@Param({"1024"})
	public int avgChunkSizeBytes;

	@Param({"25"})
	public int delayMillis;

	@Param({"10"})
	public int jitterMillis;

	private ScheduledExecutorService scheduler;
	private ChunkedBlobUploader uploader;
	private Path file;
	private Digest blobDigest;
	private Random latencyJitter;

	@Setup(Level.Trial)
	public void setup() throws Exception {
	scheduler = Executors.newScheduledThreadPool(schedulerThreads);
	latencyJitter = new Random(54321L);

	GrpcCacheClient grpcCacheClient = mock(GrpcCacheClient.class);
	CombinedCache combinedCache = mock(CombinedCache.class);

	byte[] data = new byte[fileSizeBytes];
	new Random(42).nextBytes(data);
	blobDigest = DIGEST_UTIL.compute(data);

	FileSystem fs = new InMemoryFileSystem(new JavaClock(), DigestHashFunction.SHA256);
	file = fs.getPath("/bench/blob.bin");
	file.getParentDirectory().createDirectoryAndParents();
	try (var out = file.getOutputStream()) {
	out.write(data);
	}

	FastCdcChunkingConfig chunkingConfig = new FastCdcChunkingConfig(avgChunkSizeBytes, 2, 0);
	uploader =
	new ChunkedBlobUploader(grpcCacheClient, combinedCache, chunkingConfig, DIGEST_UTIL);

	List<Digest> chunkDigests;
	try (var input = file.getInputStream()) {
	chunkDigests = new FastCdcChunker(chunkingConfig, DIGEST_UTIL).chunkToDigests(input);
	}

	when(grpcCacheClient.findMissingDigests(any(), any()))
	.thenReturn(Futures.immediateFuture(ImmutableSet.copyOf(chunkDigests)));
	when(grpcCacheClient.spliceBlob(any(), any(Digest.class), any(), any()))
	.thenReturn(Futures.immediateVoidFuture());
	when(combinedCache.uploadBlob(any(), any(Digest.class), any(Blob.class), anyBoolean()))
	.thenAnswer(
	invocation ->
	delayedFuture(null, delayMillis, jitterMillis, latencyJitter, scheduler));
	}

	@TearDown(Level.Trial)
	public void tearDown() {
	scheduler.shutdownNow();
	}
	}

	private static <T> ListenableFuture<T> delayedFuture(
	T value,
	int delayMillis,
	int jitterMillis,
	Random latencyJitter,
	ScheduledExecutorService scheduler) {
	SettableFuture<T> future = SettableFuture.create();
	var unused =
	scheduler.schedule(
	() -> future.set(value),
	jitteredDelayMillis(delayMillis, jitterMillis, latencyJitter),
	TimeUnit.MILLISECONDS);
	return future;
	}

	private static int jitteredDelayMillis(int delayMillis, int jitterMillis, Random latencyJitter) {
	if (jitterMillis == 0) {
	return delayMillis;
	}
	return Math.max(0, delayMillis + latencyJitter.nextInt((jitterMillis * 2) + 1) - jitterMillis);
	}
	}