Skip to content

Commit

Permalink
Add test
Browse files Browse the repository at this point in the history
  • Loading branch information
@coeuvre
coeuvre committed Nov 14, 2022
1 parent 27263aa commit 070f32d
Show file tree
Hide file tree
Showing 11 changed files with 335 additions and 75 deletions.
7 changes: 3 additions & 4 deletions src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.channels.ReadableByteChannel;
import java.time.Instant;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
Expand Down Expand Up @@ -157,8 +158,7 @@ void flush() throws IOException {
TreeFileArtifact.createTreeOutput(parent, parentRelativePath);
var metadata = createRemoteMetadata(remoteFile);
tree.putChild(child, metadata);
remoteLeaseService.add(
metadata.getDigest(), metadata.getSize(), metadata.getLocationIndex());
remoteLeaseService.add(metadata);
}
});

Expand All @@ -170,8 +170,7 @@ void flush() throws IOException {
if (remoteFile != null) {
var metadata = createRemoteMetadata(remoteFile);
metadataInjector.injectFile(output, metadata);
remoteLeaseService.add(
metadata.getDigest(), metadata.getSize(), metadata.getLocationIndex());
remoteLeaseService.add(metadata);
}
}
}
Expand Down
92 changes: 69 additions & 23 deletions src/main/java/com/google/devtools/build/lib/remote/RemoteLeaseService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@
import static com.google.common.base.Preconditions.checkState;

import build.bazel.remote.execution.v2.Digest;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableSet;
import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
import com.google.devtools.build.lib.events.Event;
import com.google.devtools.build.lib.events.EventHandler;
import com.google.devtools.build.lib.profiler.Profiler;
Expand All @@ -23,7 +25,7 @@
import java.time.Instant;
import java.util.Arrays;
import java.util.Collection;
import java.util.TreeSet;
import java.util.TreeMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.locks.ReentrantLock;
import javax.annotation.concurrent.GuardedBy;
Expand All @@ -43,29 +45,43 @@ public class RemoteLeaseService implements LeaseService {

private final ReentrantLock leaseLock = new ReentrantLock();

// Use Map instead of Set because we need to get the reference to the Lease and check
// `expireAtEpochMilli` when checking whether it is alive.
@GuardedBy("leaseLock")
private final TreeSet<Lease> leaseTreeSet =
new TreeSet<>(
private final TreeMap<Lease, Lease> leases =
new TreeMap<>(
(o1, o2) -> {
if (o1.equals(o2)) {
return 0;
}
return o1.expireAtEpochMilli < o2.expireAtEpochMilli ? -1 : 1;
});

private static class Lease {
private static final Lease MIN = now(Long.MIN_VALUE);
@VisibleForTesting
static class Lease {
private static final Lease MIN = now(0);

// This doesn't contribute to the equality of the lease but is used for sorting so we can
// qucikly collect a set of leases that are expired.
private final long expireAtEpochMilli;
private final byte[] digest;
private final long size;
private final int locationIndex;

public static Lease now(long now) {
private static Lease now(long now) {
return new Lease(now, new byte[0], /* size= */ 0, /* locationIndex= */ 0);
}

private Lease(long expireAtEpochMilli, byte[] digest, long size, int locationIndex) {
static Lease create(long expireAtEpochMilli, RemoteFileArtifactValue metadata) {
return new Lease(
expireAtEpochMilli,
metadata.getDigest(),
metadata.getSize(),
metadata.getLocationIndex());
}

@VisibleForTesting
Lease(long expireAtEpochMilli, byte[] digest, long size, int locationIndex) {
this.expireAtEpochMilli = expireAtEpochMilli;
this.digest = digest;
this.size = size;
Expand Down Expand Up @@ -149,13 +165,14 @@ private void readLeases() {
var leaseStore = LeaseStore.parseFrom(leasesPath.getInputStream());
leaseLock.lock();
try {
for (var lease : leaseStore.getLeasesList()) {
leaseTreeSet.add(
for (var persistedLease : leaseStore.getLeasesList()) {
var lease =
new Lease(
lease.getExpireAtEpochMilli(),
lease.getDigest().toByteArray(),
lease.getSize(),
lease.getLocationIndex()));
persistedLease.getExpireAtEpochMilli(),
persistedLease.getDigest().toByteArray(),
persistedLease.getSize(),
persistedLease.getLocationIndex());
this.leases.put(lease, lease);
}
} finally {
leaseLock.unlock();
Expand All @@ -171,7 +188,7 @@ private void writeLeases() {

leaseLock.lock();
try {
for (var lease : leaseTreeSet) {
for (var lease : this.leases.keySet()) {
leaseStore.addLeases(
RemoteLease.Lease.newBuilder()
.setExpireAtEpochMilli(lease.expireAtEpochMilli)
Expand All @@ -198,29 +215,43 @@ private Path getLeasesPath() {
return cacheDirectory.getRelative("lease_store");
}

public void add(byte[] digest, long size, int locationIndex) {
public void add(RemoteFileArtifactValue metadata) {
var now = Instant.now().toEpochMilli();
var lease = new Lease(now + remoteCacheAge.toMillis(), digest, size, locationIndex);
add(Lease.create(now + remoteCacheAge.toMillis(), metadata));
}

@VisibleForTesting
void add(Lease lease) {
leaseLock.lock();
try {
leaseTreeSet.add(lease);
leases.put(lease, lease);
} finally {
leaseLock.unlock();
}
}

@VisibleForTesting
boolean isAlive(Lease lease) {
return isAlive(lease.digest, lease.size, lease.locationIndex);
}

@Override
public boolean isAlive(byte[] digest, long size, int locationIndex) {
var now = Instant.now().toEpochMilli();
leaseLock.lock();
try {
return leaseTreeSet.contains(
new Lease(/* expireAtEpochMilli */ 0, digest, size, locationIndex));
var lease = leases.get(new Lease(/* expireAtEpochMilli */ 0, digest, size, locationIndex));
if (lease == null) {
return false;
}
return now < lease.expireAtEpochMilli;
} finally {
leaseLock.unlock();
}
}

private void renewLeases(long now, Collection<Lease> leasesToRenew)
@VisibleForTesting
void renewLeases(long now, Collection<Lease> leasesToRenew)
throws InterruptedException, ExecutionException {
if (leasesToRenew.isEmpty()) {
return;
Expand All @@ -233,9 +264,9 @@ private void renewLeases(long now, Collection<Lease> leasesToRenew)
try {
for (var lease : renewedLeases) {
if (lease.expireAtEpochMilli <= now) {
leaseTreeSet.remove(lease);
leases.remove(lease);
} else {
leaseTreeSet.add(lease);
leases.put(lease, lease);
}
}
} finally {
Expand Down Expand Up @@ -273,7 +304,7 @@ private ImmutableSet<Lease> doRenewLease(long now, Collection<Lease> leases)
private ImmutableSet<Lease> getLeasesToRenew(long now) {
leaseLock.lock();
try {
return ImmutableSet.copyOf(leaseTreeSet.subSet(Lease.MIN, Lease.now(now)));
return ImmutableSet.copyOf(leases.subMap(Lease.MIN, Lease.now(now)).keySet());
} finally {
leaseLock.unlock();
}
Expand All @@ -293,4 +324,19 @@ private void renewThreadMain() {
}
}
}

@VisibleForTesting
ImmutableSet<Lease> getAllLeases() {
leaseLock.lock();
try {
return ImmutableSet.copyOf(leases.keySet());
} finally {
leaseLock.unlock();
}
}

@VisibleForTesting
RemoteCache getRemoteCache() {
return remoteCache;
}
}
7 changes: 7 additions & 0 deletions src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.google.devtools.build.lib.util.AbruptExitException;
import com.google.devtools.build.lib.vfs.BatchStat;
import com.google.devtools.build.lib.vfs.FileSystem;
import com.google.devtools.build.lib.vfs.LeaseService;
import com.google.devtools.build.lib.vfs.ModifiedFileSet;
import com.google.devtools.build.lib.vfs.OutputService;
import com.google.devtools.build.lib.vfs.PathFragment;
Expand Down Expand Up @@ -179,4 +180,10 @@ public ArtifactPathResolver createPathResolverForArtifactValues(
remoteLeaseService);
return ArtifactPathResolver.createPathResolver(remoteFileSystem, fileSystem.getPath(execRoot));
}

@Nullable
@Override
public LeaseService getLeaseService() {
return remoteLeaseService;
}
}
21 changes: 14 additions & 7 deletions src/test/java/com/google/devtools/build/lib/actions/ActionCacheCheckerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,8 @@ private void runAction(
metadataHandler,
/*artifactExpander=*/ null,
platform,
/*isRemoteCacheEnabled=*/ true);
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null);
if (token != null) {
// Real action execution would happen here.
ActionExecutionContext context = mock(ActionExecutionContext.class);
Expand Down Expand Up @@ -443,7 +444,8 @@ public void testDeletedConstantMetadataOutputCausesReexecution() throws Exceptio
new FakeMetadataHandler(),
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ true))
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null))
.isNotNull();
}

Expand Down Expand Up @@ -569,7 +571,8 @@ public void saveOutputMetadata_remoteFileMetadataLoaded() throws Exception {
metadataHandler,
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ true);
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null);

assertThat(output.getPath().exists()).isFalse();
assertThat(token).isNull();
Expand Down Expand Up @@ -598,7 +601,8 @@ public void saveOutputMetadata_remoteOutputUnavailable_remoteFileMetadataNotLoad
metadataHandler,
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ false);
/*isRemoteCacheEnabled=*/ false,
/*leaseService=*/ null);

assertThat(output.getPath().exists()).isFalse();
assertThat(token).isNotNull();
Expand Down Expand Up @@ -772,7 +776,8 @@ public void saveOutputMetadata_emptyTreeMetadata_notSaved() throws Exception {
metadataHandler,
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ true);
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null);

assertThat(token).isNull();
assertThat(output.getPath().exists()).isFalse();
Expand Down Expand Up @@ -870,7 +875,8 @@ public void saveOutputMetadata_treeMetadata_remoteFileMetadataLoaded() throws Ex
metadataHandler,
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ true);
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null);

TreeArtifactValue expectedMetadata =
createTreeMetadata(
Expand Down Expand Up @@ -1031,7 +1037,8 @@ public void saveOutputMetadata_treeMetadataWithSameLocalFileMetadata_cached(
metadataHandler,
/*artifactExpander=*/ null,
/*remoteDefaultPlatformProperties=*/ ImmutableMap.of(),
/*isRemoteCacheEnabled=*/ true);
/*isRemoteCacheEnabled=*/ true,
/*leaseService=*/ null);

assertThat(token).isNull();
assertStatistics(1, new MissDetailsBuilder().set(MissReason.NOT_CACHED, 1).build());
Expand Down
15 changes: 9 additions & 6 deletions ...t/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploaderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,8 @@ public void remoteFileShouldNotBeUploaded_actionFs() throws Exception {
outputRoot.getRoot().asPath().relativeTo(execRoot).getPathString(),
outputs,
ImmutableList.of(artifact),
actionInputFetcher, remoteLeaseService);
actionInputFetcher,
mock(RemoteLeaseService.class));
Path remotePath = remoteFs.getPath(artifact.getPath().getPathString());
assertThat(remotePath.getFileSystem()).isEqualTo(remoteFs);
LocalFile file =
Expand Down Expand Up @@ -428,7 +429,7 @@ public void remoteFileShouldNotBeUploaded_findMissingDigests() throws Exception
PathConverter pathConverter = artifactUploader.upload(files).get();

// assert
verify(digestQuerier).findMissingDigests(any(), any());
verify(digestQuerier).findMissingDigests(any(), any(), any());
verify(remoteCache).uploadFile(any(), eq(localDigest), any());
assertThat(pathConverter.apply(remoteFile)).contains(remoteDigest.getHash());
assertThat(pathConverter.apply(localFile)).contains(localDigest.getHash());
Expand Down Expand Up @@ -468,9 +469,11 @@ private RemoteCache newRemoteCache(
doAnswer(
invocationOnMock ->
missingDigestsFinder.findMissingDigests(
invocationOnMock.getArgument(0), invocationOnMock.getArgument(1)))
invocationOnMock.getArgument(0),
invocationOnMock.getArgument(1),
invocationOnMock.getArgument(2)))
.when(cacheClient)
.findMissingDigests(any(), any());
.findMissingDigests(any(), any(), any());

return new RemoteCache(
CacheCapabilities.getDefaultInstance(), cacheClient, remoteOptions, DIGEST_UTIL);
Expand Down Expand Up @@ -500,7 +503,7 @@ public StaticMissingDigestsFinder(ImmutableSet<Digest> knownDigests) {

@Override
public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(
RemoteActionExecutionContext context, Iterable<Digest> digests) {
RemoteActionExecutionContext context, Intention intention, Iterable<Digest> digests) {
ImmutableSet.Builder<Digest> missingDigests = ImmutableSet.builder();
for (Digest digest : digests) {
if (!knownDigests.contains(digest)) {
Expand All @@ -517,7 +520,7 @@ private static class AllMissingDigestsFinder implements MissingDigestsFinder {

@Override
public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(
RemoteActionExecutionContext context, Iterable<Digest> digests) {
RemoteActionExecutionContext context, Intention intention, Iterable<Digest> digests) {
return Futures.immediateFuture(ImmutableSet.copyOf(digests));
}
}
Expand Down
3 changes: 2 additions & 1 deletion src/test/java/com/google/devtools/build/lib/remote/RemoteActionFileSystemTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,8 @@ protected RemoteActionFileSystem createActionFileSystem(
outputRoot.getRoot().asPath().relativeTo(execRoot).getPathString(),
inputs,
outputs,
inputFetcher, remoteLeaseService);
inputFetcher,
mock(RemoteLeaseService.class));
remoteActionFileSystem.updateContext(metadataInjector);
remoteActionFileSystem.createDirectoryAndParents(outputRoot.getRoot().asPath().asFragment());
return remoteActionFileSystem;
Expand Down
Loading

0 comments on commit 070f32d

Please sign in to comment.