-
Notifications
You must be signed in to change notification settings - Fork 4.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add (inactive) code to compress large GenQuery results in memory on t…
…he fly Dynamically switches to gzipping GenQuery results >1M, trading CPU for memory. Hard coded to off for now while I sort out how to connect this to a flag. PiperOrigin-RevId: 258568242
- Loading branch information
1 parent
26bcd27
commit 6d63f8a
Showing
3 changed files
with
392 additions
and
15 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
210 changes: 210 additions & 0 deletions
210
src/main/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStream.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,210 @@ | ||
| // Copyright 2019 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.rules.genquery; | ||
|
|
||
| import com.google.common.annotations.VisibleForTesting; | ||
| import com.google.common.base.Preconditions; | ||
| import com.google.common.io.ByteStreams; | ||
| import com.google.devtools.build.lib.util.Fingerprint; | ||
| import com.google.protobuf.ByteString; | ||
| import java.io.IOException; | ||
| import java.io.OutputStream; | ||
| import java.util.zip.GZIPInputStream; | ||
| import java.util.zip.GZIPOutputStream; | ||
|
|
||
| /** | ||
| * {@link OutputStream} implementation optimized for {@link GenQuery} by (optionally) compressing | ||
| * query results on the fly. Produces {@link GenQueryResult}s which are preferred for storing the | ||
| * output of {@link GenQuery}'s underlying queries. | ||
| */ | ||
| class GenQueryOutputStream extends OutputStream { | ||
|
|
||
| /** | ||
| * When compression is enabled, the threshold at which the stream will switch to compressing | ||
| * output. The value of this constant is arbitrary but effective. | ||
| */ | ||
| private static final int COMPRESSION_THRESHOLD = 1 << 20; | ||
|
|
||
| /** | ||
| * Encapsulates the output of a {@link GenQuery}'s query. CPU and memory overhead of individual | ||
| * methods depends on the underlying content and settings. | ||
| */ | ||
| interface GenQueryResult { | ||
| /** Returns the query output as a {@link ByteString}. */ | ||
| ByteString getBytes() throws IOException; | ||
|
|
||
| /** | ||
| * Adds the query output to the supplied {@link Fingerprint}. Equivalent to {@code | ||
| * fingerprint.addBytes(genQueryResult.getBytes())}, but potentially more efficient. | ||
| */ | ||
| void fingerprint(Fingerprint fingerprint); | ||
|
|
||
| /** | ||
| * Returns the size of the output. This must be a constant-time operation for all | ||
| * implementations. | ||
| */ | ||
| int size(); | ||
|
|
||
| /** | ||
| * Writes the query output to the provided {@link OutputStream}. Equivalent to {@code | ||
| * genQueryResult.getBytes().writeTo(out)}, but potentially more efficient. | ||
| */ | ||
| void writeTo(OutputStream out) throws IOException; | ||
| } | ||
|
|
||
| private final boolean compressionEnabled; | ||
| private int bytesWritten = 0; | ||
| private boolean compressed = false; | ||
| private boolean closed = false; | ||
| private ByteString.Output bytesOut = ByteString.newOutput(); | ||
| private OutputStream out = bytesOut; | ||
|
|
||
| GenQueryOutputStream(boolean compressionEnabled) { | ||
| this.compressionEnabled = compressionEnabled; | ||
| } | ||
|
|
||
| @Override | ||
| public void write(int b) throws IOException { | ||
| maybeStartCompression(1); | ||
| out.write(b); | ||
| bytesWritten += 1; | ||
| } | ||
|
|
||
| @Override | ||
| public void write(byte[] bytes) throws IOException { | ||
| write(bytes, 0, bytes.length); | ||
| } | ||
|
|
||
| @Override | ||
| public void write(byte[] bytes, int off, int len) throws IOException { | ||
| maybeStartCompression(len); | ||
| out.write(bytes, off, len); | ||
| bytesWritten += len; | ||
| } | ||
|
|
||
| @Override | ||
| public void flush() throws IOException { | ||
| out.flush(); | ||
| } | ||
|
|
||
| @Override | ||
| public void close() throws IOException { | ||
| out.close(); | ||
| closed = true; | ||
| } | ||
|
|
||
| GenQueryResult getResult() { | ||
| Preconditions.checkState(closed, "Must be closed"); | ||
| return compressed | ||
| ? new CompressedResult(bytesOut.toByteString(), bytesWritten) | ||
| : new RegularResult(bytesOut.toByteString()); | ||
| } | ||
|
|
||
| private void maybeStartCompression(int additionalBytes) throws IOException { | ||
| if (!compressionEnabled) { | ||
| return; | ||
| } | ||
|
|
||
| if (compressed) { | ||
| return; | ||
| } | ||
|
|
||
| if (bytesWritten + additionalBytes < COMPRESSION_THRESHOLD) { | ||
| return; | ||
| } | ||
|
|
||
| ByteString.Output compressedBytesOut = ByteString.newOutput(); | ||
| GZIPOutputStream gzipOut = new GZIPOutputStream(compressedBytesOut); | ||
| bytesOut.writeTo(gzipOut); | ||
| bytesOut = compressedBytesOut; | ||
| out = gzipOut; | ||
| compressed = true; | ||
| } | ||
|
|
||
| @VisibleForTesting | ||
| static class RegularResult implements GenQueryResult { | ||
| private final ByteString data; | ||
|
|
||
| RegularResult(ByteString data) { | ||
| this.data = data; | ||
| } | ||
|
|
||
| @Override | ||
| public ByteString getBytes() { | ||
| return data; | ||
| } | ||
|
|
||
| @Override | ||
| public int size() { | ||
| return data.size(); | ||
| } | ||
|
|
||
| @Override | ||
| public void fingerprint(Fingerprint fingerprint) { | ||
| fingerprint.addBytes(data); | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(OutputStream out) throws IOException { | ||
| data.writeTo(out); | ||
| } | ||
| } | ||
|
|
||
| @VisibleForTesting | ||
| static class CompressedResult implements GenQueryResult { | ||
| private final ByteString compressedData; | ||
| private final int size; | ||
|
|
||
| CompressedResult(ByteString compressedData, int size) { | ||
| this.compressedData = compressedData; | ||
| this.size = size; | ||
| } | ||
|
|
||
| @Override | ||
| public ByteString getBytes() throws IOException { | ||
| ByteString.Output out = ByteString.newOutput(size); | ||
| try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) { | ||
| ByteStreams.copy(gzipIn, out); | ||
| } | ||
| return out.toByteString(); | ||
| } | ||
|
|
||
| @Override | ||
| public int size() { | ||
| return size; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(OutputStream out) throws IOException { | ||
| try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) { | ||
| ByteStreams.copy(gzipIn, out); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void fingerprint(Fingerprint fingerprint) { | ||
| try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) { | ||
| byte[] chunk = new byte[4092]; | ||
| int bytesRead; | ||
| while ((bytesRead = gzipIn.read(chunk)) > 0) { | ||
| fingerprint.addBytes(chunk, 0, bytesRead); | ||
| } | ||
| } catch (IOException e) { | ||
| // Unexpected, everything should be in memory! | ||
| throw new IllegalStateException("Unexpected IOException", e); | ||
| } | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.