// 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.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionInputFileCache;
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
import com.google.devtools.build.lib.remote.RemoteProtocol.BlobChunk;
import com.google.devtools.build.lib.remote.RemoteProtocol.ContentDigest;
import com.google.devtools.build.lib.util.Preconditions;
import com.google.devtools.build.lib.vfs.Path;
import com.google.protobuf.ByteString;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Set;
import javax.annotation.Nullable;
/** An iterator-type object that transforms byte sources into a stream of BlobChunk messages. */
public final class Chunker {
/** An Item is an opaque digestable source of bytes. */
interface Item {
ContentDigest getDigest() throws IOException;
InputStream getInputStream() throws IOException;
}
private final Iterator<Item> inputIterator;
private InputStream currentStream;
private final Set<ContentDigest> digests;
private ContentDigest digest;
private long bytesLeft;
private final int chunkSize;
Chunker(
Iterator<Item> inputIterator,
int chunkSize,
// If present, specifies which digests to output out of the whole input.
@Nullable Set<ContentDigest> digests)
throws IOException {
Preconditions.checkArgument(chunkSize > 0, "Chunk size must be greater than 0");
this.digests = digests;
this.inputIterator = inputIterator;
this.chunkSize = chunkSize;
advanceInput();
}
Chunker(Iterator<Item> inputIterator, int chunkSize) throws IOException {
this(inputIterator, chunkSize, null);
}
Chunker(Item input, int chunkSize) throws IOException {
this(Iterators.singletonIterator(input), chunkSize, ImmutableSet.of(input.getDigest()));
}
private void advanceInput() throws IOException {
do {
if (inputIterator != null && inputIterator.hasNext()) {
Item input = inputIterator.next();
digest = input.getDigest();
currentStream = input.getInputStream();
bytesLeft = digest.getSizeBytes();
} else {
digest = null;
currentStream = null;
bytesLeft = 0;
}
} while (digest != null && digests != null && !digests.contains(digest));
}
/** True if the object has more BlobChunk elements. */
public boolean hasNext() {
return currentStream != null;
}
/** Consume the next BlobChunk element. */
public BlobChunk next() throws IOException {
if (!hasNext()) {
throw new NoSuchElementException();
}
BlobChunk.Builder chunk = BlobChunk.newBuilder();
long offset = digest.getSizeBytes() - bytesLeft;
if (offset == 0) {
chunk.setDigest(digest);
} else {
chunk.setOffset(offset);
}
if (bytesLeft > 0) {
byte[] blob = new byte[(int) Math.min(bytesLeft, chunkSize)];
currentStream.read(blob);
chunk.setData(ByteString.copyFrom(blob));
bytesLeft -= blob.length;
}
if (bytesLeft == 0) {
currentStream.close();
advanceInput(); // Sets the current stream to null, if it was the last.
}
return chunk.build();
}
static Item toItem(final byte[] blob) {
return new Item() {
@Override
public ContentDigest getDigest() throws IOException {
return ContentDigests.computeDigest(blob);
}
@Override
public InputStream getInputStream() throws IOException {
return new ByteArrayInputStream(blob);
}
};
}
static Item toItem(final Path file) {
return new Item() {
@Override
public ContentDigest getDigest() throws IOException {
return ContentDigests.computeDigest(file);
}
@Override
public InputStream getInputStream() throws IOException {
return file.getInputStream();
}
};
}
static Item toItem(
final ActionInput input, final ActionInputFileCache inputCache, final Path execRoot) {
if (input instanceof VirtualActionInput) {
return toItem((VirtualActionInput) input);
}
return new Item() {
@Override
public ContentDigest getDigest() throws IOException {
return ContentDigests.getDigestFromInputCache(input, inputCache);
}
@Override
public InputStream getInputStream() throws IOException {
return execRoot.getRelative(input.getExecPathString()).getInputStream();
}
};
}
static Item toItem(final VirtualActionInput input) {
return new Item() {
@Override
public ContentDigest getDigest() throws IOException {
return ContentDigests.computeDigest(input);
}
@Override
public InputStream getInputStream() throws IOException {
ByteArrayOutputStream buffer = new ByteArrayOutputStream();
input.writeTo(buffer);
return new ByteArrayInputStream(buffer.toByteArray());
}
};
}
/**
* Create a Chunker from a given ActionInput, taking its digest from the provided
* ActionInputFileCache.
*/
public static Chunker from(
ActionInput input, int chunkSize, ActionInputFileCache inputCache, Path execRoot)
throws IOException {
return new Chunker(toItem(input, inputCache, execRoot), chunkSize);
}
/** Create a Chunker from a given blob and chunkSize. */
public static Chunker from(byte[] blob, int chunkSize) throws IOException {
return new Chunker(toItem(blob), chunkSize);
}
/** Create a Chunker from a given Path and chunkSize. */
public static Chunker from(Path file, int chunkSize) throws IOException {
return new Chunker(toItem(file), chunkSize);
}
/**
* Create a Chunker from multiple input sources. The order of the sources provided to the Builder
* will be the same order they will be chunked by.
*/
public static final class Builder {
private final ArrayList<Item> items = new ArrayList<>();
private Set<ContentDigest> digests = null;
private int chunkSize = 0;
public Chunker build() throws IOException {
return new Chunker(items.iterator(), chunkSize, digests);
}
public Builder chunkSize(int chunkSize) {
this.chunkSize = chunkSize;
return this;
}
/**
* Restricts the Chunker to use only inputs with these digests. This is an optimization for CAS
* uploads where a list of digests missing from the CAS is known.
*/
public Builder onlyUseDigests(Set<ContentDigest> digests) {
this.digests = digests;
return this;
}
public Builder addInput(byte[] blob) {
items.add(toItem(blob));
return this;
}
public Builder addInput(Path file) {
items.add(toItem(file));
return this;
}
public Builder addInput(ActionInput input, ActionInputFileCache inputCache, Path execRoot) {
items.add(toItem(input, inputCache, execRoot));
return this;
}
public Builder addAllInputs(
Collection<? extends ActionInput> inputs, ActionInputFileCache inputCache, Path execRoot) {
for (ActionInput input : inputs) {
items.add(toItem(input, inputCache, execRoot));
}
return this;
}
}
}