ExternalizableBean: do not retain oversized thread-local buffer after write

dump/src/util/dump/ExternalizableBean.java

@@ -815,6 +815,7 @@ default int readExternalFieldNumber( ObjectInput in ) throws IOException {
 
    @Override
    default void writeExternal( ObjectOutput out ) throws IOException {
+      StreamCache streamCache = null;
       try {
          ClassConfig _config = getConfig(getClass());
 
@@ -826,7 +827,6 @@ default void writeExternal( ObjectOutput out ) throws IOException {
          }
 
          ObjectOutput originalOut = out;
-         StreamCache streamCache = null;
          BytesCache bytesCache = null;
          ObjectOutput cachingOut = null;
 
@@ -1207,13 +1207,24 @@ default void writeExternal( ObjectOutput out ) throws IOException {
                originalOut.writeInt(bytesCache.size());
                bytesCache.writeTo(originalOut);
                out = originalOut;
+               bytesCache.reset();
                streamCache._inUse = false;
             }
          }
       }
       catch ( Exception e ) {
          throw new RuntimeException("Failed to externalize class " + getClass().getName(), e);
       }
+      finally {
+         if ( streamCache != null && streamCache._inUse ) {
+            // a field failed to write: the cached ObjectOutput may hold partial state, so never reuse this StreamCache.
+            // Replace it in the thread-local (if it is the one from there), which also drops its buffer.
+            streamCache._bytesCache.reset();
+            if ( STREAM_CACHE.get() == streamCache ) {
+               STREAM_CACHE.set(new StreamCache());
+            }
+         }
+      }
    }
 
    /**

dump/src/util/dump/ExternalizationHelper.java

@@ -58,6 +58,10 @@ class ExternalizationHelper {
    static final Map<Class, Boolean>      CLASS_CHANGED_INCOMPATIBLY = new ConcurrentHashMap<>();
    static final ThreadLocal<StreamCache> STREAM_CACHE               = ThreadLocal.withInitial(StreamCache::new);
 
+   /** Buffers larger than this are not retained by a {@link BytesCache} after a write finished; they are replaced by a small one. */
+   static final int BYTES_CACHE_MAX_RETAINED_CAPACITY = 1024 * 1024;
+   static final int BYTES_CACHE_INITIAL_CAPACITY      = 1024;
+
    static {
       try {
          boolean config = Boolean.parseBoolean(System.getProperty("ExternalizableBean.USE_UNSAFE_FIELD_ACCESSORS", "true"));
@@ -1075,7 +1079,7 @@ static class BytesCache extends OutputStream {
       protected int    _count;
 
       public BytesCache() {
-         this(1024);
+         this(BYTES_CACHE_INITIAL_CAPACITY);
       }
 
       public BytesCache( int size ) {
@@ -1085,11 +1089,14 @@ public BytesCache( int size ) {
          _buffer = new byte[size];
       }
 
+      /**
+       * Marks the buffer as empty and drops an oversized backing array, so that it is not retained (e.g. by a long-living thread).
+       */
       public void reset() {
          _count = 0;
-         if ( _buffer.length > 1048576 ) {
+         if ( _buffer.length > BYTES_CACHE_MAX_RETAINED_CAPACITY ) {
             // let it shrink
-            _buffer = new byte[1024];
+            _buffer = new byte[BYTES_CACHE_INITIAL_CAPACITY];
          }
       }
 

dump/test/util/dump/ExternalizationBufferRetentionTest.java

@@ -0,0 +1,147 @@
+package util.dump;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import org.junit.Test;
+
+import util.dump.ExternalizableBean.externalize;
+import util.dump.ExternalizationHelper.BytesCache;
+import util.dump.ExternalizationHelper.StreamCache;
+
+public class ExternalizationBufferRetentionTest {
+
+   public static class LargeBean implements ExternalizableBean {
+
+      // List fields have a dynamic length and therefore go through the thread-local buffer.
+      // The large field must be written last: a following dynamic field would reset() (and thereby shrink) the buffer anyway.
+      @externalize(1)
+      public List<LargeBean> _nested;
+
+      @externalize(2)
+      public List<byte[]> _data;
+
+      public LargeBean() {}
+
+      public LargeBean( byte[] data, LargeBean nested ) {
+         _data = new ArrayList<>(Collections.singletonList(data));
+         _nested = nested == null ? new ArrayList<>() : new ArrayList<>(Collections.singletonList(nested));
+      }
+   }
+
+   private static byte[] write( LargeBean bean ) throws Exception {
+      ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+      try ( ObjectOutputStream out = new ObjectOutputStream(bytes) ) {
+         bean.writeExternal(out);
+      }
+      return bytes.toByteArray();
+   }
+
+   @Test
+   public void largeWriteDoesNotRetainLargeBuffer() throws Exception {
+      int size = ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY * 4;
+      byte[] serialized = write(new LargeBean(new byte[size], null));
+      assertThat(serialized.length).isGreaterThan(size);
+
+      StreamCache cache = ExternalizationHelper.STREAM_CACHE.get();
+      assertThat(cache._inUse).isFalse();
+      assertThat(cache._bytesCache._buffer.length).isLessThanOrEqualTo(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY);
+   }
+
+   @Test
+   public void smallWriteKeepsBufferForReuse() throws Exception {
+      write(new LargeBean(new byte[10_000], null));
+      StreamCache cache = ExternalizationHelper.STREAM_CACHE.get();
+      byte[] retained = cache._bytesCache._buffer;
+      assertThat(retained.length).isGreaterThanOrEqualTo(10_000);
+
+      write(new LargeBean(new byte[10_000], null));
+      assertThat(cache._bytesCache._buffer).isSameAs(retained);
+   }
+
+   @Test
+   public void nestedLargeWriteRoundtripsAndDoesNotRetain() throws Exception {
+      int size = ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY * 2;
+      LargeBean inner = new LargeBean(new byte[size], null);
+      LargeBean outer = new LargeBean(new byte[size], inner);
+      byte[] serialized = write(outer);
+
+      LargeBean read = new LargeBean();
+      try ( ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(serialized)) ) {
+         read.readExternal(in);
+      }
+      assertThat(read._data.get(0)).hasSize(size);
+      assertThat(read._nested.get(0)._data.get(0)).hasSize(size);
+
+      StreamCache cache = ExternalizationHelper.STREAM_CACHE.get();
+      assertThat(cache._inUse).isFalse();
+      assertThat(cache._bytesCache._buffer.length).isLessThanOrEqualTo(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY);
+   }
+
+   @Test
+   public void failedWriteDoesNotPoisonNextWrite() throws Exception {
+      LargeBean broken = new LargeBean(new byte[ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY * 2], null);
+      broken._data.add(null);
+      broken._nested.add(new LargeBean(null, null) {
+         @Override
+         public void writeExternal( java.io.ObjectOutput out ) {
+            throw new IllegalStateException("boom");
+         }
+      });
+      try {
+         write(broken);
+         org.junit.Assert.fail("expected failure");
+      }
+      catch ( RuntimeException expected ) {
+         // expected
+      }
+
+      StreamCache cache = ExternalizationHelper.STREAM_CACHE.get();
+      assertThat(cache._inUse).isFalse();
+      assertThat(cache._bytesCache._buffer.length).isLessThanOrEqualTo(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY);
+
+      LargeBean good = new LargeBean(new byte[] { 1, 2, 3 }, null);
+      LargeBean read = new LargeBean();
+      try ( ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(write(good))) ) {
+         read.readExternal(in);
+      }
+      assertThat(read._data.get(0)).containsExactly(1, 2, 3);
+   }
+
+   @Test
+   public void largeWriteOnWorkerThreadDoesNotRetainLargeBufferOnThatThread() throws Exception {
+      // the buffer is retained per thread, so check the thread-local of the long-living thread that wrote
+      java.util.concurrent.ExecutorService executor = java.util.concurrent.Executors.newSingleThreadExecutor();
+      try {
+         int size = ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY * 4;
+         executor.submit(() -> write(new LargeBean(new byte[size], null))).get();
+         int retainedOnWorker = executor.submit(() -> ExternalizationHelper.STREAM_CACHE.get()._bytesCache._buffer.length).get();
+         assertThat(retainedOnWorker).isLessThanOrEqualTo(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY);
+      }
+      finally {
+         executor.shutdownNow();
+      }
+   }
+
+   @Test
+   public void bytesCacheKeepsBufferAtCapAndDropsAboveCap() {
+      BytesCache atCap = new BytesCache(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY);
+      byte[] kept = atCap._buffer;
+      atCap.reset();
+      assertThat(atCap._buffer).isSameAs(kept);
+      assertThat(atCap.size()).isZero();
+
+      BytesCache aboveCap = new BytesCache(ExternalizationHelper.BYTES_CACHE_MAX_RETAINED_CAPACITY + 1);
+      aboveCap.reset();
+      assertThat(aboveCap._buffer.length).isEqualTo(ExternalizationHelper.BYTES_CACHE_INITIAL_CAPACITY);
+      assertThat(aboveCap.size()).isZero();
+   }
+}