| // Copyright 2024 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.skyframe.serialization; |
| |
| import static com.google.common.util.concurrent.MoreExecutors.directExecutor; |
| import static com.google.devtools.build.lib.skyframe.serialization.FutureHelpers.aggregateStatusFutures; |
| import static com.google.devtools.build.lib.skyframe.serialization.FutureHelpers.getDoneSerializationFuture; |
| import static com.google.devtools.build.lib.skyframe.serialization.FutureHelpers.reportAnyFailures; |
| import static com.google.devtools.build.lib.skyframe.serialization.FutureHelpers.waitForSerializationFuture; |
| |
| import com.google.common.annotations.VisibleForTesting; |
| import com.google.common.collect.ImmutableClassToInstanceMap; |
| import com.google.common.collect.Lists; |
| import com.google.common.util.concurrent.Futures; |
| import com.google.common.util.concurrent.ListenableFuture; |
| import com.google.common.util.concurrent.SettableFuture; |
| import com.google.protobuf.ByteString; |
| import com.google.protobuf.CodedOutputStream; |
| import java.io.ByteArrayOutputStream; |
| import java.io.IOException; |
| import java.util.ArrayList; |
| import java.util.List; |
| import javax.annotation.Nullable; |
| |
| /** |
| * A {@link SerializationContext} that supports both memoization and shared subobjects. |
| * |
| * <p>Sharing subobjects occurs by means of uploading them asynchronously to a {@link |
| * #fingerprintValueService} store. The status of these uploads may be observed through {@link |
| * #createFutureToBlockWritingOn} and {@link SerializationResult#getFutureToBlockWritesOn}. |
| */ |
| final class SharedValueSerializationContext extends MemoizingSerializationContext { |
| private final FingerprintValueService fingerprintValueService; |
| |
| /** |
| * Futures that represent writes to remote storage. |
| * |
| * <p>For consistency, the serialized bytes should not be published for other consumers until |
| * these writes complete. |
| */ |
| @Nullable // lazily initialized |
| private ArrayList<ListenableFuture<Void>> futuresToBlockWritingOn; |
| |
| /** |
| * Futures that mark deferred data bytes when using {@link #putSharedValue}. |
| * |
| * <p>The serialized byte representation contains placeholders until these futures complete. |
| */ |
| @Nullable // lazily initialized |
| private ArrayList<FuturePut> futurePuts; |
| |
| @VisibleForTesting // private |
| static SharedValueSerializationContext createForTesting( |
| ObjectCodecRegistry codecRegistry, |
| ImmutableClassToInstanceMap<Object> dependencies, |
| FingerprintValueService fingerprintValueService) { |
| return new SharedValueSerializationContext( |
| codecRegistry, dependencies, fingerprintValueService); |
| } |
| |
| private SharedValueSerializationContext( |
| ObjectCodecRegistry codecRegistry, |
| ImmutableClassToInstanceMap<Object> dependencies, |
| FingerprintValueService fingerprintValueService) { |
| super(codecRegistry, dependencies); |
| this.fingerprintValueService = fingerprintValueService; |
| } |
| |
| /** |
| * Serializes {@code subject} and returns a result that may have an associated future. |
| * |
| * <p>This method may block when the serialization of a shared child is performed by another |
| * thread. It blocks until the fingerprints of such shared children become available. However, |
| * this method does not block on uploads to the {@link FingerprintValueService}. The status of the |
| * upload is provided by {@link SerializationResult#getFutureToBlockWritesOn}. |
| */ |
| static SerializationResult<ByteString> serializeToResult( |
| ObjectCodecRegistry codecRegistry, |
| ImmutableClassToInstanceMap<Object> dependencies, |
| FingerprintValueService fingerprintValueService, |
| @Nullable Object subject) |
| throws SerializationException { |
| ByteArrayOutputStream bytesOut = new ByteArrayOutputStream(); |
| CodedOutputStream codedOut = CodedOutputStream.newInstance(bytesOut); |
| SharedValueSerializationContext context = |
| new SharedValueSerializationContext(codecRegistry, dependencies, fingerprintValueService); |
| try { |
| context.serialize(subject, codedOut); |
| codedOut.flush(); |
| } catch (IOException e) { |
| throw new SerializationException("Failed to serialize: " + subject, e); |
| } |
| return context.createResult(bytesOut.toByteArray()); |
| } |
| |
| @Override |
| public <T> void putSharedValue( |
| T child, |
| @Nullable Object distinguisher, |
| DeferredObjectCodec<T> codec, |
| CodedOutputStream codedOut) |
| throws IOException, SerializationException { |
| SettableFuture<PutOperation> putOperation = SettableFuture.create(); |
| Object previous = |
| fingerprintValueService.getOrClaimPutOperation(child, distinguisher, putOperation); |
| if (previous != null) { |
| // Uses the previous result and discards `putOperation`. |
| if (previous instanceof ListenableFuture) { |
| // HINT: this is the only place where `FuturePut`s originate. The other caller of |
| // `recordFuturePut` propagates them transitively. |
| @SuppressWarnings("unchecked") |
| ListenableFuture<PutOperation> inflight = (ListenableFuture<PutOperation>) previous; |
| recordFuturePut(inflight, codedOut); |
| return; |
| } |
| codedOut.writeRawBytes((ByteString) previous); |
| return; |
| } |
| |
| try { |
| putOwnedSharedValue(child, codec, codedOut, putOperation); |
| } catch (SerializationException | IOException e) { |
| putOperation.setException(e); |
| throw e; |
| } catch (RuntimeException | Error e) { |
| // `putOperation` must always be set to avoid deadlock-like behaviors in its consumers. |
| putOperation.setException(e); |
| throw e; |
| } |
| } |
| |
| /** |
| * Implementation of {@link #putSharedValue} when it owns the {@code putOperation}. |
| * |
| * <p>This method does the following. |
| * |
| * <ul> |
| * <li>Uploads {@code child}'s bytes into {@link #fingerprintValueService}. |
| * <li>Sets the result of the upload in {@code putOperation}. |
| * <li>Writes the fingerprint of {@code child}'s serialized bytes (or a placeholder) to {@code |
| * codedOut}. |
| * <li>Adds either a {@link #futuresToBlockWritingOn} entry on {@link #futurePuts} entry. |
| * </ul> |
| * |
| * <p>There are two cases, in particular. |
| * |
| * <ul> |
| * <li><em>Serializing {@code child} causes recursive calls to {@link #putSharedValue} and some |
| * of those recursive calls are are still being processed in a different thread</em>: this |
| * results in placeholders being added to the {@code codedOut} and {@link #futurePuts} |
| * entries to be recorded. |
| * <li><em>Any other case, for example, serializing {@code child} does not recursively call |
| * {@link #putSharedValue} or if it does, the fingerprints are either already available or |
| * owned by this thread</em>: here, the serialized bytes of {@code child} are computed in |
| * the calling thread and become immediately available. The status of uploading the |
| * serialized {@code child} is added to {@link #futuresToBlockWritingOn}. |
| * </ul> |
| * |
| * <p>When entries are added to {@link #futurePuts}, the caller must ensure that for every entry, |
| * {@link FuturePut#fillInPlaceholderAndReturnWriteStatus} is called to update placeholders |
| * present in the serialized bytes. |
| */ |
| private <T> void putOwnedSharedValue( |
| T child, |
| DeferredObjectCodec<T> codec, |
| CodedOutputStream codedOut, |
| SettableFuture<PutOperation> putOperation) |
| throws IOException, SerializationException { |
| byte[] childBytes; // bytes of serialized child, maybe containing placeholders |
| // This is initially populated with the `futuresToWriteBlockingOn` of the child. If the child |
| // has `FuturePut`s, the write statuses of those puts will be added to this list. |
| ArrayList<ListenableFuture<Void>> childWriteStatuses; |
| // Each entry of the following list corresponds to a `putSharedValue` call made by serialization |
| // of `child`. Used to fill-in `childBytes` placeholders. |
| ArrayList<FuturePut> childDeferredBytes; |
| { // Serializes `child` with a fresh context, populating the above variables. |
| ByteArrayOutputStream childStream = new ByteArrayOutputStream(); |
| CodedOutputStream childCodedOut = CodedOutputStream.newInstance(childStream); |
| SharedValueSerializationContext childContext = getFreshContext(); |
| codec.serialize(childContext, child, childCodedOut); |
| childCodedOut.flush(); |
| |
| childBytes = childStream.toByteArray(); |
| childWriteStatuses = |
| childContext.futuresToBlockWritingOn == null |
| ? new ArrayList<>() |
| : childContext.futuresToBlockWritingOn; |
| childDeferredBytes = childContext.futurePuts; |
| } |
| |
| if (childDeferredBytes == null) { |
| // There are no deferred bytes so `childBytes` is complete. Starts the upload. |
| ByteString fingerprint = fingerprintValueService.fingerprint(childBytes); |
| codedOut.writeRawBytes(fingerprint); // Writes only the fingerprint to the stream. |
| childWriteStatuses.add(fingerprintValueService.put(fingerprint, childBytes)); |
| |
| ListenableFuture<Void> writeStatus = aggregateStatusFutures(childWriteStatuses); |
| putOperation.set(PutOperation.create(fingerprint, writeStatus)); |
| addFutureToBlockWritingOn(writeStatus); |
| return; |
| } |
| |
| // Ensures enough space for the subvalues and one additional for the put of `child` itself. |
| childWriteStatuses.ensureCapacity(childWriteStatuses.size() + childDeferredBytes.size() + 1); |
| |
| // There are pending fingerprint sub-computations. |
| ListenableFuture<PutOperation> futureStore = |
| whenAllPutsStarted(childDeferredBytes) |
| .call( |
| () -> { |
| fillPlaceholdersAndCollectWriteStatuses( |
| childDeferredBytes, childBytes, childWriteStatuses); |
| // All placeholders are filled-in. Starts the upload. |
| ByteString fingerprint = fingerprintValueService.fingerprint(childBytes); |
| childWriteStatuses.add(fingerprintValueService.put(fingerprint, childBytes)); |
| return PutOperation.create( |
| fingerprint, aggregateStatusFutures(childWriteStatuses)); |
| }, |
| directExecutor()); |
| |
| putOperation.setFuture(futureStore); |
| recordFuturePut(futureStore, codedOut); |
| } |
| |
| @Override |
| public SharedValueSerializationContext getFreshContext() { |
| return new SharedValueSerializationContext( |
| getCodecRegistry(), getDependencies(), fingerprintValueService); |
| } |
| |
| @Override |
| public void addFutureToBlockWritingOn(ListenableFuture<Void> future) { |
| if (futuresToBlockWritingOn == null) { |
| futuresToBlockWritingOn = new ArrayList<>(); |
| } |
| futuresToBlockWritingOn.add(future); |
| } |
| |
| @Override |
| @Nullable |
| public ListenableFuture<Void> createFutureToBlockWritingOn() { |
| if (futurePuts == null) { |
| if (futuresToBlockWritingOn == null) { |
| return null; |
| } |
| return aggregateStatusFutures(futuresToBlockWritingOn); |
| } |
| |
| List<ListenableFuture<Void>> writeStatuses = initializeCombinedWriteStatusesList(); |
| for (FuturePut futurePut : futurePuts) { |
| writeStatuses.add( |
| Futures.transformAsync( |
| futurePut.getFuturePut(), PutOperation::writeStatus, directExecutor())); |
| } |
| return aggregateStatusFutures(writeStatuses); |
| } |
| |
| /** |
| * Records a deferred fingerprint value to be written to the serialized bytes. |
| * |
| * <p>Adds the corresponding entry to {@link #futurePuts} and inserts placeholder bytes into |
| * {@code codedOut}. |
| */ |
| private void recordFuturePut(ListenableFuture<PutOperation> futurePut, CodedOutputStream codedOut) |
| throws IOException { |
| if (futurePuts == null) { |
| futurePuts = new ArrayList<>(); |
| } |
| futurePuts.add(new FuturePut(codedOut.getTotalBytesWritten(), futurePut)); |
| // Adds a placeholder for the real fingerprint to be filled in after the future completes. |
| codedOut.writeRawBytes(fingerprintValueService.fingerprintPlaceholder()); |
| } |
| |
| private SerializationResult<ByteString> createResult(byte[] bytes) throws SerializationException { |
| // TODO: b/297857068 - If ByteString.copyFrom overhead is excessive, use reflection to avoid it. |
| if (futurePuts == null) { |
| ByteString finalBytes = ByteString.copyFrom(bytes); |
| return futuresToBlockWritingOn == null |
| ? SerializationResult.createWithoutFuture(finalBytes) |
| : SerializationResult.create(finalBytes, aggregateStatusFutures(futuresToBlockWritingOn)); |
| } |
| |
| List<ListenableFuture<Void>> writeStatuses = initializeCombinedWriteStatusesList(); |
| // `futurePuts` are produced when multiple threads race to write the same value. This is |
| // expected to be relatively rare and waiting here is for CPU bound serialization work only. |
| // |
| // TODO: b/297857068 - It's possible to expose this intermediate future to callers as a |
| // SerializationResult<ListenableFuture<byte[]>>. Do so if the contention outweighs the extra |
| // complexity. |
| try { |
| waitForSerializationFuture( |
| whenAllPutsStarted(futurePuts) |
| .call( |
| () -> { |
| fillPlaceholdersAndCollectWriteStatuses(futurePuts, bytes, writeStatuses); |
| return null; |
| }, |
| directExecutor())); |
| } catch (SerializationException e) { |
| // An exception has occurred. Ensures that any additional errors from writing are handled |
| // before throwing the exception. |
| reportAnyFailures(Futures.whenAllSucceed(futuresToBlockWritingOn)); |
| throw e; |
| } |
| return SerializationResult.create( |
| ByteString.copyFrom(bytes), aggregateStatusFutures(writeStatuses)); |
| } |
| |
| /** |
| * Initializes a write status list including any {@link #futuresToBlockWritingOn} and sized for |
| * the caller to add statuses from {@link #futurePuts}. |
| * |
| * <p>The caller must ensure that {@link #futurePuts} is non-null. |
| */ |
| private List<ListenableFuture<Void>> initializeCombinedWriteStatusesList() { |
| if (futuresToBlockWritingOn == null) { |
| return new ArrayList<>(futurePuts.size()); |
| } |
| ArrayList<ListenableFuture<Void>> statuses = |
| new ArrayList<>(futurePuts.size() + futuresToBlockWritingOn.size()); |
| statuses.addAll(futuresToBlockWritingOn); |
| return statuses; |
| } |
| |
| private static Futures.FutureCombiner<?> whenAllPutsStarted(List<FuturePut> stores) { |
| return Futures.whenAllSucceed(Lists.transform(stores, FuturePut::getFuturePut)); |
| } |
| |
| /** |
| * Fills placeholders in {@code bytes} from {@code futurePuts} and adds write statuses to {@code |
| * writeStatuses}. |
| * |
| * <p>The caller must ensure that all the {@link FuturePut#futurePut}s futures in {@code |
| * futurePuts} are done before calling this, that is, {@link #whenAllPutsStarted}. |
| */ |
| private static void fillPlaceholdersAndCollectWriteStatuses( |
| List<FuturePut> futurePuts, byte[] bytes, List<ListenableFuture<Void>> writeStatuses) |
| throws SerializationException { |
| for (FuturePut put : futurePuts) { |
| writeStatuses.add(put.fillInPlaceholderAndReturnWriteStatus(bytes)); |
| } |
| } |
| |
| /** |
| * A tuple consisting of a stream offset and a pending {@link PutOperation}. |
| * |
| * <p>When {@link PutOperation} becomes available, {@link fillInPlaceholderAndReturnWriteStatus} |
| * must be called. |
| */ |
| private static class FuturePut { |
| /** Where the fingerprint should be written when it becomes available. */ |
| private final int offset; |
| |
| /** |
| * A {@link PutOperation} that may have not yet started. |
| * |
| * <p>This arises when associated bytes could be mid-computation in a different thread. The |
| * future completes when those bytes are known and its asynchronous write to the {@link |
| * #fingerprintValueService} has started. {@link PutOperation#writeStatus} provides the status |
| * of that write. |
| */ |
| private final ListenableFuture<PutOperation> futurePut; |
| |
| private FuturePut(int offset, ListenableFuture<PutOperation> futurePut) { |
| this.offset = offset; |
| this.futurePut = futurePut; |
| } |
| |
| private ListenableFuture<PutOperation> getFuturePut() { |
| return futurePut; |
| } |
| |
| /** |
| * Copies the fingerprint into {@code bytes} at {@link #offset}. |
| * |
| * <p>The caller must ensure that {@link #futurePut} is complete before calling this method. |
| */ |
| private ListenableFuture<Void> fillInPlaceholderAndReturnWriteStatus(byte[] bytes) |
| throws SerializationException { |
| PutOperation put = getDoneSerializationFuture(futurePut); |
| put.fingerprint().copyTo(bytes, offset); |
| return put.writeStatus(); |
| } |
| } |
| } |