src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreActionCache.java - bazel - Git at Google

 // Copyright 2016 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 com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.remote.TreeNodeRepository.TreeNode;
 import com.google.devtools.build.lib.remote.blobstore.SimpleBlobStore;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.DigestUtil.ActionKey;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.remoteexecution.v1test.ActionResult;
 import com.google.devtools.remoteexecution.v1test.Command;
 import com.google.devtools.remoteexecution.v1test.Digest;
 import com.google.devtools.remoteexecution.v1test.Directory;
 import com.google.devtools.remoteexecution.v1test.DirectoryNode;
 import com.google.devtools.remoteexecution.v1test.FileNode;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.Collection;
 import java.util.Map;

 /**
  * A RemoteActionCache implementation that uses a concurrent map as a distributed storage for files
  * and action output.
  *
  * <p>The thread safety is guaranteed by the underlying map.
  *
  * <p>Note that this class is used from src/tools/remote.
  */
 @ThreadSafe
 public final class SimpleBlobStoreActionCache extends AbstractRemoteActionCache {
   private static final int MAX_BLOB_SIZE_FOR_INLINE = 10 * 1024;

   private final SimpleBlobStore blobStore;

   public SimpleBlobStoreActionCache(
       RemoteOptions options, SimpleBlobStore blobStore, DigestUtil digestUtil) {
     super(options, digestUtil);
     this.blobStore = blobStore;
   }

   @Override
   public void ensureInputsPresent(
       TreeNodeRepository repository, Path execRoot, TreeNode root, Command command)
           throws IOException, InterruptedException {
     repository.computeMerkleDigests(root);
     uploadBlob(command.toByteArray());
     for (Directory directory : repository.treeToDirectories(root)) {
       uploadBlob(directory.toByteArray());
     }
     // TODO(ulfjack): Only upload files that aren't in the CAS yet?
     for (TreeNode leaf : repository.leaves(root)) {
       uploadFileContents(leaf.getActionInput(), execRoot, repository.getInputFileCache());
     }
   }

   public void downloadTree(Digest rootDigest, Path rootLocation)
       throws IOException, InterruptedException {
     FileSystemUtils.createDirectoryAndParents(rootLocation);
     Directory directory = Directory.parseFrom(downloadBlob(rootDigest));
     for (FileNode file : directory.getFilesList()) {
       downloadFile(
           rootLocation.getRelative(file.getName()), file.getDigest(), file.getIsExecutable(), null);
     }
     for (DirectoryNode child : directory.getDirectoriesList()) {
       downloadTree(child.getDigest(), rootLocation.getRelative(child.getName()));
     }
   }

   private Digest uploadFileContents(Path file) throws IOException, InterruptedException {
     Digest digest = digestUtil.compute(file);
     try (InputStream in = file.getInputStream()) {
       return uploadStream(digest, in);
     }
   }

   private Digest uploadFileContents(
       ActionInput input, Path execRoot, MetadataProvider inputCache)
           throws IOException, InterruptedException {
     if (input instanceof VirtualActionInput) {
       byte[] blob = ((VirtualActionInput) input).getBytes().toByteArray();
       return uploadBlob(blob, digestUtil.compute(blob));
     }
     try (InputStream in = execRoot.getRelative(input.getExecPathString()).getInputStream()) {
       return uploadStream(DigestUtil.getFromInputCache(input, inputCache), in);
     }
   }

   @Override
   public void upload(
       ActionKey actionKey,
       Path execRoot,
       Collection<Path> files,
       FileOutErr outErr,
       boolean uploadAction)
       throws ExecException, IOException, InterruptedException {
     ActionResult.Builder result = ActionResult.newBuilder();
     upload(result, execRoot, files);
     if (outErr.getErrorPath().exists()) {
       Digest stderr = uploadFileContents(outErr.getErrorPath());
       result.setStderrDigest(stderr);
     }
     if (outErr.getOutputPath().exists()) {
       Digest stdout = uploadFileContents(outErr.getOutputPath());
       result.setStdoutDigest(stdout);
     }
     if (uploadAction) {
       blobStore.putActionResult(actionKey.getDigest().getHash(), result.build().toByteArray());
     }
   }

   public void upload(ActionResult.Builder result, Path execRoot, Collection<Path> files)
       throws ExecException, IOException, InterruptedException {
     UploadManifest manifest =
         new UploadManifest(digestUtil, result, execRoot, options.allowSymlinkUpload);
     manifest.addFiles(files);

     for (Map.Entry<Digest, Path> entry : manifest.getDigestToFile().entrySet()) {
       try (InputStream in = entry.getValue().getInputStream()) {
         uploadStream(entry.getKey(), in);
       }
     }

     for (Map.Entry<Digest, Chunker> entry : manifest.getDigestToChunkers().entrySet()) {
       uploadBlob(entry.getValue().next().getData().toByteArray(), entry.getKey());
     }
   }

   public void uploadOutErr(ActionResult.Builder result, byte[] stdout, byte[] stderr)
       throws IOException, InterruptedException {
     if (stdout.length <= MAX_BLOB_SIZE_FOR_INLINE) {
       result.setStdoutRaw(ByteString.copyFrom(stdout));
     } else if (stdout.length > 0) {
       result.setStdoutDigest(uploadBlob(stdout));
     }
     if (stderr.length <= MAX_BLOB_SIZE_FOR_INLINE) {
       result.setStderrRaw(ByteString.copyFrom(stderr));
     } else if (stderr.length > 0) {
       result.setStderrDigest(uploadBlob(stderr));
     }
   }

   public Digest uploadBlob(byte[] blob) throws IOException, InterruptedException {
     return uploadBlob(blob, digestUtil.compute(blob));
   }

   private Digest uploadBlob(byte[] blob, Digest digest) throws IOException, InterruptedException {
     try (InputStream in = new ByteArrayInputStream(blob)) {
       return uploadStream(digest, in);
     }
   }

   public Digest uploadStream(Digest digest, InputStream in)
       throws IOException, InterruptedException {
     blobStore.put(digest.getHash(), digest.getSizeBytes(), in);
     return digest;
   }

   public boolean containsKey(Digest digest) throws IOException, InterruptedException {
     return blobStore.containsKey(digest.getHash());
   }

   @Override
   public ActionResult getCachedActionResult(ActionKey actionKey)
       throws IOException, InterruptedException {
     try {
       byte[] data = downloadActionResult(actionKey.getDigest());
       return ActionResult.parseFrom(data);
     } catch (InvalidProtocolBufferException | CacheNotFoundException e) {
       return null;
     }
   }

   private byte[] downloadActionResult(Digest digest) throws IOException, InterruptedException {
     if (digest.getSizeBytes() == 0) {
       return new byte[0];
     }
     // This unconditionally downloads the whole blob into memory!
     ByteArrayOutputStream out = new ByteArrayOutputStream();
     boolean success = blobStore.getActionResult(digest.getHash(), out);
     if (!success) {
       throw new CacheNotFoundException(digest, digestUtil);
     }
     return out.toByteArray();
   }

   public void setCachedActionResult(ActionKey actionKey, ActionResult result)
       throws IOException, InterruptedException {
     blobStore.putActionResult(actionKey.getDigest().getHash(), result.toByteArray());
   }

   @Override
   public void close() {
     blobStore.close();
   }

   @Override
   protected void downloadBlob(Digest digest, Path dest) throws IOException, InterruptedException {
     try (OutputStream out = dest.getOutputStream()) {
       boolean success = blobStore.get(digest.getHash(), out);
       if (!success) {
         throw new CacheNotFoundException(digest, digestUtil);
       }
     }
   }

   @Override
   public byte[] downloadBlob(Digest digest) throws IOException, InterruptedException {
     if (digest.getSizeBytes() == 0) {
       return new byte[0];
     }
     try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
       boolean success = blobStore.get(digest.getHash(), out);
       if (!success) {
         throw new CacheNotFoundException(digest, digestUtil);
       }
       return out.toByteArray();
     }
   }
 }
	// Copyright 2016 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 com.google.devtools.build.lib.actions.ActionInput;
	import com.google.devtools.build.lib.actions.ExecException;
	import com.google.devtools.build.lib.actions.MetadataProvider;
	import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
	import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
	import com.google.devtools.build.lib.remote.TreeNodeRepository.TreeNode;
	import com.google.devtools.build.lib.remote.blobstore.SimpleBlobStore;
	import com.google.devtools.build.lib.remote.util.DigestUtil;
	import com.google.devtools.build.lib.remote.util.DigestUtil.ActionKey;
	import com.google.devtools.build.lib.util.io.FileOutErr;
	import com.google.devtools.build.lib.vfs.FileSystemUtils;
	import com.google.devtools.build.lib.vfs.Path;
	import com.google.devtools.remoteexecution.v1test.ActionResult;
	import com.google.devtools.remoteexecution.v1test.Command;
	import com.google.devtools.remoteexecution.v1test.Digest;
	import com.google.devtools.remoteexecution.v1test.Directory;
	import com.google.devtools.remoteexecution.v1test.DirectoryNode;
	import com.google.devtools.remoteexecution.v1test.FileNode;
	import com.google.protobuf.ByteString;
	import com.google.protobuf.InvalidProtocolBufferException;
	import java.io.ByteArrayInputStream;
	import java.io.ByteArrayOutputStream;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.OutputStream;
	import java.util.Collection;
	import java.util.Map;

	/**
	* A RemoteActionCache implementation that uses a concurrent map as a distributed storage for files
	* and action output.
	*
	* <p>The thread safety is guaranteed by the underlying map.
	*
	* <p>Note that this class is used from src/tools/remote.
	*/
	@ThreadSafe
	public final class SimpleBlobStoreActionCache extends AbstractRemoteActionCache {
	private static final int MAX_BLOB_SIZE_FOR_INLINE = 10 * 1024;

	private final SimpleBlobStore blobStore;

	public SimpleBlobStoreActionCache(
	RemoteOptions options, SimpleBlobStore blobStore, DigestUtil digestUtil) {
	super(options, digestUtil);
	this.blobStore = blobStore;
	}

	@Override
	public void ensureInputsPresent(
	TreeNodeRepository repository, Path execRoot, TreeNode root, Command command)
	throws IOException, InterruptedException {
	repository.computeMerkleDigests(root);
	uploadBlob(command.toByteArray());
	for (Directory directory : repository.treeToDirectories(root)) {
	uploadBlob(directory.toByteArray());
	}
	// TODO(ulfjack): Only upload files that aren't in the CAS yet?
	for (TreeNode leaf : repository.leaves(root)) {
	uploadFileContents(leaf.getActionInput(), execRoot, repository.getInputFileCache());
	}
	}

	public void downloadTree(Digest rootDigest, Path rootLocation)
	throws IOException, InterruptedException {
	FileSystemUtils.createDirectoryAndParents(rootLocation);
	Directory directory = Directory.parseFrom(downloadBlob(rootDigest));
	for (FileNode file : directory.getFilesList()) {
	downloadFile(
	rootLocation.getRelative(file.getName()), file.getDigest(), file.getIsExecutable(), null);
	}
	for (DirectoryNode child : directory.getDirectoriesList()) {
	downloadTree(child.getDigest(), rootLocation.getRelative(child.getName()));
	}
	}

	private Digest uploadFileContents(Path file) throws IOException, InterruptedException {
	Digest digest = digestUtil.compute(file);
	try (InputStream in = file.getInputStream()) {
	return uploadStream(digest, in);
	}
	}

	private Digest uploadFileContents(
	ActionInput input, Path execRoot, MetadataProvider inputCache)
	throws IOException, InterruptedException {
	if (input instanceof VirtualActionInput) {
	byte[] blob = ((VirtualActionInput) input).getBytes().toByteArray();
	return uploadBlob(blob, digestUtil.compute(blob));
	}
	try (InputStream in = execRoot.getRelative(input.getExecPathString()).getInputStream()) {
	return uploadStream(DigestUtil.getFromInputCache(input, inputCache), in);
	}
	}

	@Override
	public void upload(
	ActionKey actionKey,
	Path execRoot,
	Collection<Path> files,
	FileOutErr outErr,
	boolean uploadAction)
	throws ExecException, IOException, InterruptedException {
	ActionResult.Builder result = ActionResult.newBuilder();
	upload(result, execRoot, files);
	if (outErr.getErrorPath().exists()) {
	Digest stderr = uploadFileContents(outErr.getErrorPath());
	result.setStderrDigest(stderr);
	}
	if (outErr.getOutputPath().exists()) {
	Digest stdout = uploadFileContents(outErr.getOutputPath());
	result.setStdoutDigest(stdout);
	}
	if (uploadAction) {
	blobStore.putActionResult(actionKey.getDigest().getHash(), result.build().toByteArray());
	}
	}

	public void upload(ActionResult.Builder result, Path execRoot, Collection<Path> files)
	throws ExecException, IOException, InterruptedException {
	UploadManifest manifest =
	new UploadManifest(digestUtil, result, execRoot, options.allowSymlinkUpload);
	manifest.addFiles(files);

	for (Map.Entry<Digest, Path> entry : manifest.getDigestToFile().entrySet()) {
	try (InputStream in = entry.getValue().getInputStream()) {
	uploadStream(entry.getKey(), in);
	}
	}

	for (Map.Entry<Digest, Chunker> entry : manifest.getDigestToChunkers().entrySet()) {
	uploadBlob(entry.getValue().next().getData().toByteArray(), entry.getKey());
	}
	}

	public void uploadOutErr(ActionResult.Builder result, byte[] stdout, byte[] stderr)
	throws IOException, InterruptedException {
	if (stdout.length <= MAX_BLOB_SIZE_FOR_INLINE) {
	result.setStdoutRaw(ByteString.copyFrom(stdout));
	} else if (stdout.length > 0) {
	result.setStdoutDigest(uploadBlob(stdout));
	}
	if (stderr.length <= MAX_BLOB_SIZE_FOR_INLINE) {
	result.setStderrRaw(ByteString.copyFrom(stderr));
	} else if (stderr.length > 0) {
	result.setStderrDigest(uploadBlob(stderr));
	}
	}

	public Digest uploadBlob(byte[] blob) throws IOException, InterruptedException {
	return uploadBlob(blob, digestUtil.compute(blob));
	}

	private Digest uploadBlob(byte[] blob, Digest digest) throws IOException, InterruptedException {
	try (InputStream in = new ByteArrayInputStream(blob)) {
	return uploadStream(digest, in);
	}
	}

	public Digest uploadStream(Digest digest, InputStream in)
	throws IOException, InterruptedException {
	blobStore.put(digest.getHash(), digest.getSizeBytes(), in);
	return digest;
	}

	public boolean containsKey(Digest digest) throws IOException, InterruptedException {
	return blobStore.containsKey(digest.getHash());
	}

	@Override
	public ActionResult getCachedActionResult(ActionKey actionKey)
	throws IOException, InterruptedException {
	try {
	byte[] data = downloadActionResult(actionKey.getDigest());
	return ActionResult.parseFrom(data);
	} catch (InvalidProtocolBufferException \| CacheNotFoundException e) {
	return null;
	}
	}

	private byte[] downloadActionResult(Digest digest) throws IOException, InterruptedException {
	if (digest.getSizeBytes() == 0) {
	return new byte[0];
	}
	// This unconditionally downloads the whole blob into memory!
	ByteArrayOutputStream out = new ByteArrayOutputStream();
	boolean success = blobStore.getActionResult(digest.getHash(), out);
	if (!success) {
	throw new CacheNotFoundException(digest, digestUtil);
	}
	return out.toByteArray();
	}

	public void setCachedActionResult(ActionKey actionKey, ActionResult result)
	throws IOException, InterruptedException {
	blobStore.putActionResult(actionKey.getDigest().getHash(), result.toByteArray());
	}

	@Override
	public void close() {
	blobStore.close();
	}

	@Override
	protected void downloadBlob(Digest digest, Path dest) throws IOException, InterruptedException {
	try (OutputStream out = dest.getOutputStream()) {
	boolean success = blobStore.get(digest.getHash(), out);
	if (!success) {
	throw new CacheNotFoundException(digest, digestUtil);
	}
	}
	}

	@Override
	public byte[] downloadBlob(Digest digest) throws IOException, InterruptedException {
	if (digest.getSizeBytes() == 0) {
	return new byte[0];
	}
	try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
	boolean success = blobStore.get(digest.getHash(), out);
	if (!success) {
	throw new CacheNotFoundException(digest, digestUtil);
	}
	return out.toByteArray();
	}
	}
	}