From 4ea2de33a5741aa6c31233ba3cf4f95250f04553 Mon Sep 17 00:00:00 2001 From: Oli Gillespie Date: Fri, 3 Nov 2023 15:07:34 +0000 Subject: [PATCH 1/2] Allow AccessibleByteArrayOutputStream to shrink AccessibleByteArrayOutputStream currently can only grow or stay the same size. This could be a problem if it is used for a rare large item, but then used only for small items - the buffer would stay large forever, consuming memory. Since an application may hold many of these buffers via their use particularly in Cipher, the memory impact could be significant. This simple shrinkage algorithm shrinks the buffer if it has been 'too large' N times in a row. N starts at 1, and grows exponentially up to a limit of 1024 as re-growths are triggered. N starts low to enable reacting quickly, but as more growths are observed it increases to avoid churn. It is capped to avoid degrading into 'never-shrink' behaviour over the long term. Other ideas were considered, like keeping track of the last N used sizes and making decisions based on those values, but this algorithm has a tiny footprint (2 ints, could be shorts) and low computation overhead. --- .../AccessibleByteArrayOutputStream.java | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java b/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java index 0a4124ba..316a91e0 100644 --- a/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java +++ b/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java @@ -7,6 +7,10 @@ import java.util.Arrays; class AccessibleByteArrayOutputStream extends OutputStream implements Cloneable { + private static final int MAX_OVERSIZED_THRESHOLD = 1024; + private int timesOversized; + private int timesOversizedThreshold; + private final int limit; private byte[] buf; private int count; @@ -83,11 +87,22 @@ byte[] getDataBuffer() { } void reset() { - Arrays.fill(buf, 0, count, (byte) 0); + int sizeUsed = count; + Arrays.fill(buf, 0, sizeUsed, (byte) 0); count = 0; - // TODO: Consider keeping track of length at reset. - // If it is consistently below the maximum value we may want to trim - // down to save on memory. + + // Consider shrinking the buffer. + if (sizeUsed * 2 < buf.length) { + // The buffer was over-sized for this usage. + if (timesOversized++ > timesOversizedThreshold) { + // Shrink the buffer. + buf = new byte[buf.length / 2]; + timesOversized = 0; + } + } else { + // Buffer was not over-sized, reset counter. + timesOversized = 0; + } } void write(final ByteBuffer bbuff) { @@ -122,5 +137,8 @@ private void growCapacity(final int newCapacity, final boolean doNotAllocateMore final byte[] toZeroize = buf; buf = tmp; Arrays.fill(toZeroize, 0, count, (byte) 0); + + // Every time we need to grow, make it harder to shrink in the future (up to a limit). + timesOversizedThreshold = Math.min(MAX_OVERSIZED_THRESHOLD, timesOversized * 2); } } From d23cfe4c08e1ddca67bdba0b6eb921517fbfe1fb Mon Sep 17 00:00:00 2001 From: Oli Gillespie Date: Thu, 9 Nov 2023 11:19:01 +0000 Subject: [PATCH 2/2] Allow different buffer shrink strategy implementations --- .../AccessibleByteArrayOutputStream.java | 27 ++--- .../crypto/provider/BufferShrinkStrategy.java | 103 ++++++++++++++++++ 2 files changed, 112 insertions(+), 18 deletions(-) create mode 100644 src/com/amazon/corretto/crypto/provider/BufferShrinkStrategy.java diff --git a/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java b/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java index 316a91e0..61243af8 100644 --- a/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java +++ b/src/com/amazon/corretto/crypto/provider/AccessibleByteArrayOutputStream.java @@ -7,11 +7,8 @@ import java.util.Arrays; class AccessibleByteArrayOutputStream extends OutputStream implements Cloneable { - private static final int MAX_OVERSIZED_THRESHOLD = 1024; - private int timesOversized; - private int timesOversizedThreshold; - private final int limit; + private final BufferShrinkStrategy shrinkStrategy; private byte[] buf; private int count; @@ -23,6 +20,10 @@ class AccessibleByteArrayOutputStream extends OutputStream implements Cloneable } AccessibleByteArrayOutputStream(final int capacity, final int limit) { + this(capacity, limit, new BufferShrinkStrategy.BasicThreshold()); + } + + AccessibleByteArrayOutputStream(final int capacity, final int limit, BufferShrinkStrategy shrinkStrategy) { if (limit < 0) { throw new IllegalArgumentException("Limit must be non-negative"); } @@ -31,6 +32,7 @@ class AccessibleByteArrayOutputStream extends OutputStream implements Cloneable } buf = capacity == 0 ? Utils.EMPTY_ARRAY : new byte[capacity]; this.limit = limit; + this.shrinkStrategy = shrinkStrategy; count = 0; } @@ -91,17 +93,9 @@ void reset() { Arrays.fill(buf, 0, sizeUsed, (byte) 0); count = 0; - // Consider shrinking the buffer. - if (sizeUsed * 2 < buf.length) { - // The buffer was over-sized for this usage. - if (timesOversized++ > timesOversizedThreshold) { - // Shrink the buffer. - buf = new byte[buf.length / 2]; - timesOversized = 0; - } - } else { - // Buffer was not over-sized, reset counter. - timesOversized = 0; + if (shrinkStrategy.shouldShrink(sizeUsed, buf.length)) { + // Shrink the buffer. + buf = new byte[buf.length / 2]; } } @@ -137,8 +131,5 @@ private void growCapacity(final int newCapacity, final boolean doNotAllocateMore final byte[] toZeroize = buf; buf = tmp; Arrays.fill(toZeroize, 0, count, (byte) 0); - - // Every time we need to grow, make it harder to shrink in the future (up to a limit). - timesOversizedThreshold = Math.min(MAX_OVERSIZED_THRESHOLD, timesOversized * 2); } } diff --git a/src/com/amazon/corretto/crypto/provider/BufferShrinkStrategy.java b/src/com/amazon/corretto/crypto/provider/BufferShrinkStrategy.java new file mode 100644 index 00000000..aae3059c --- /dev/null +++ b/src/com/amazon/corretto/crypto/provider/BufferShrinkStrategy.java @@ -0,0 +1,103 @@ +package com.amazon.corretto.crypto.provider; + +/** + * Represents a strategy for downsizing/shrinking a reusable buffer. + * For example, {@link AccessibleByteArrayOutputStream} holds a reusable buffer which grows as needed + * to accommodate different sized payloads. If this buffer never shrinks, it may waste space (for example + * if we see one rare very large payload but subsequent payloads are small, the buffer will remain unnecessarily + * large). + * Note: implementations are usually stateful and thus instances cannot be safely shared. + */ +public interface BufferShrinkStrategy { + // TODO: could return 'recommended buffer size' instead of boolean. Value/simplicity? + // TODO: should the strategy also handle growth? + + /** + * Buffer owners should call this after consumption of every payload. + * E.g. handlePayload(byte[] payload) { + * ... maybe grow buffer + * ... encrypt + * if (shouldShrink(payload.length, buffer.length)) shrinkBuf(); + * } + * @param payloadSize The size of the payload processed. + * @param bufferSize The size of the buffer. + * @return true if the strategy recommends shrinking the buffer. false otherwise. + */ + boolean shouldShrink(int payloadSize, int bufferSize); + + /** + * Shrink the buffer when it is too large for the payload ('over-sized') N times in a row. + * E.g. if an 800KB payload grows the buffer to 1MB, we shrink the buffer after seeing N consecutive + * payloads under 500KB. + */ + class BasicThreshold implements BufferShrinkStrategy { + private final int timesOversizedThreshold; + private int timesOversized = 0; + + public BasicThreshold() { + this.timesOversizedThreshold = 1024; + } + + public BasicThreshold(int timesOversizedThreshold) { + this.timesOversizedThreshold = timesOversizedThreshold; + } + + @Override + public boolean shouldShrink(int payloadSize, int bufferSize) { + if (bufferSize / 2 > payloadSize) { + // The buffer was over-sized for this usage. + if (timesOversized++ > timesOversizedThreshold) { + timesOversized = 0; + return true; + } + } else { + // Buffer was not over-sized, reset counter. + timesOversized = 0; + } + return false; + } + } + + /** + * Similar to {@link BasicThreshold}, but the threshold starts at 1 and increases + * to the chosen limit. This has the benefit of being somewhat adaptive; it starts + * out eager to shrink quickly after a large payload, but slows down every time + * re-growth is needed, up to the chosen limit. + */ + class IncreasingThreshold implements BufferShrinkStrategy { + private final int maxOversizedThreshold; + private int timesOversizedThreshold = 1; + private int timesOversized = 0; + private int previousBufferSize = Integer.MAX_VALUE; + + public IncreasingThreshold() { + this.maxOversizedThreshold = 1024; + } + + public IncreasingThreshold(int maxOversizedThreshold) { + this.maxOversizedThreshold = maxOversizedThreshold; + } + + @Override + public boolean shouldShrink(int payloadSize, int bufferSize) { + if (bufferSize > previousBufferSize) { + // Every time we need to grow, make it harder to shrink in the future (up to a limit). + timesOversizedThreshold = Math.min(maxOversizedThreshold, timesOversizedThreshold * 2); + } + previousBufferSize = bufferSize; + + if (bufferSize / 2 > payloadSize) { + // The buffer was over-sized for this usage. + if (timesOversized++ > timesOversizedThreshold) { + timesOversized = 0; + return true; + } + } else { + // Buffer was not over-sized, reset counter. + timesOversized = 0; + } + + return false; + } + } +}