KAFKA-10500: Thread Cache Resizes

streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java

@@ -160,6 +160,7 @@ public class KafkaStreams implements AutoCloseable {
     private final StreamsMetricsImpl streamsMetrics;
     private final ProcessorTopology taskTopology;
     private final ProcessorTopology globalTaskTopology;
+    private final long totalCacheSize;
 
     GlobalStreamThread globalStreamThread;
     private KafkaStreams.StateListener stateListener;
@@ -821,12 +822,8 @@ private KafkaStreams(final InternalTopologyBuilder internalTopologyBuilder,
                 "must subscribe to at least one source topic or global table.");
         }
 
-        long totalCacheSize = config.getLong(StreamsConfig.CACHE_MAX_BYTES_BUFFERING_CONFIG);
-        if (totalCacheSize < 0) {
-            totalCacheSize = 0;
-            log.warn("Negative cache size passed in. Reverting to cache size of 0 bytes.");
-        }
-        final long cacheSizePerThread = totalCacheSize / (numStreamThreads + (hasGlobalTopology ? 1 : 0));
+        totalCacheSize = config.getLong(StreamsConfig.CACHE_MAX_BYTES_BUFFERING_CONFIG);
+        final long cacheSizePerThread = getCacheSizePerThread(numStreamThreads);
         final boolean hasPersistentStores = taskTopology.hasPersistentLocalStore() ||
                 (hasGlobalTopology && globalTaskTopology.hasPersistentGlobalStore());
 
@@ -902,6 +899,17 @@ private KafkaStreams(final InternalTopologyBuilder internalTopologyBuilder,
         rocksDBMetricsRecordingService = maybeCreateRocksDBMetricsRecordingService(clientId, config);
     }
 
+    private long getCacheSizePerThread(final int numStreamThreads) {
+        return totalCacheSize / (numStreamThreads + ((globalTaskTopology != null) ? 1 : 0));
+    }
+
+    private void resizeThreadCache(final int numStreamThreads) {
+        final long cacheSizePreThread = getCacheSizePerThread(numStreamThreads);
+        for (final StreamThread streamThread: threads) {
+            streamThread.resizeCache(cacheSizePreThread);
+        }
+    }
+
     private ScheduledExecutorService setupStateDirCleaner() {
         return Executors.newSingleThreadScheduledExecutor(r -> {
             final Thread thread = new Thread(r, clientId + "-CleanupThread");

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -282,6 +282,7 @@ public boolean isRunning() {
     private final Consumer<byte[], byte[]> restoreConsumer;
     private final Admin adminClient;
     private final InternalTopologyBuilder builder;
+    private final java.util.function.Consumer<Long> cacheResizer;
 
     private java.util.function.Consumer<Throwable> streamsUncaughtExceptionHandler;
     private Runnable shutdownErrorHook;
@@ -395,7 +396,8 @@ public static StreamThread create(final InternalTopologyBuilder builder,
             referenceContainer.assignmentErrorCode,
             referenceContainer.nextScheduledRebalanceMs,
             shutdownErrorHook,
-            streamsUncaughtExceptionHandler
+            streamsUncaughtExceptionHandler,
+            cacheSize -> cache.resize(cacheSize)
         );
 
         taskManager.setPartitionResetter(partitions -> streamThread.resetOffsets(partitions, null));
@@ -448,7 +450,8 @@ public StreamThread(final Time time,
                         final AtomicInteger assignmentErrorCode,
                         final AtomicLong nextProbingRebalanceMs,
                         final Runnable shutdownErrorHook,
-                        final java.util.function.Consumer<Throwable> streamsUncaughtExceptionHandler) {
+                        final java.util.function.Consumer<Throwable> streamsUncaughtExceptionHandler,
+                        final java.util.function.Consumer<Long> cacheResizer) {
         super(threadId);
         this.stateLock = new Object();
 
@@ -468,6 +471,8 @@ public StreamThread(final Time time,
         this.assignmentErrorCode = assignmentErrorCode;
         this.shutdownErrorHook = shutdownErrorHook;
         this.streamsUncaughtExceptionHandler = streamsUncaughtExceptionHandler;
+        this.cacheResizer = cacheResizer;
+
 
         // The following sensors are created here but their references are not stored in this object, since within
         // this object they are not recorded. The sensors are created here so that the stream threads starts with all
@@ -615,6 +620,10 @@ private void subscribeConsumer() {
         }
     }
 
+    public void resizeCache(final long size) {
+        cacheResizer.accept(size);
+    }
+
     /**
      * One iteration of a thread includes the following steps:
      *

streams/src/main/java/org/apache/kafka/streams/state/internals/ThreadCache.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.streams.state.internals;
 
 import org.apache.kafka.common.utils.Bytes;
+import org.apache.kafka.common.utils.CircularIterator;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
@@ -35,7 +36,7 @@
  */
 public class ThreadCache {
     private final Logger log;
-    private final long maxCacheSizeBytes;
+    private volatile long maxCacheSizeBytes;
     private final StreamsMetricsImpl metrics;
     private final Map<String, NamedCache> caches = new HashMap<>();
 
@@ -71,6 +72,22 @@ public long flushes() {
         return numFlushes;
     }
 
+    public void resize(final long newCacheSizeBytes) {
+        final boolean shrink = newCacheSizeBytes < maxCacheSizeBytes;
+        maxCacheSizeBytes = newCacheSizeBytes;
+        if (shrink) {
+            final CircularIterator<NamedCache> circularIterator = new CircularIterator<>(caches.values());
+            while (sizeBytes() > maxCacheSizeBytes) {
+                final NamedCache cache = circularIterator.next();
+                if (cache.isEmpty()) {
+                    circularIterator.remove();
+                }
+                cache.evict();
+                numEvicts++;
+            }
+        }
+    }
+
     /**
      * The thread cache maintains a set of {@link NamedCache}s whose names are a concatenation of the task ID and the
      * underlying store name. This method creates those names.

streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java

@@ -237,7 +237,7 @@ private StreamThread createStreamThread(@SuppressWarnings("SameParameterValue")
             new MockStateRestoreListener(),
             threadIdx,
             null,
-            null
+            HANDLER
         );
     }
 
@@ -488,6 +488,7 @@ public void shouldNotCommitBeforeTheCommitInterval() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
+            HANDLER,
             null
         );
         thread.setNow(mockTime.milliseconds());
@@ -727,6 +728,7 @@ public void shouldNotCauseExceptionIfNothingCommitted() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
+            HANDLER,
             null
         );
         thread.setNow(mockTime.milliseconds());
@@ -767,6 +769,7 @@ public void shouldCommitAfterTheCommitInterval() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
+            HANDLER,
             null
         );
 
@@ -960,7 +963,8 @@ public void shouldShutdownTaskManagerOnClose() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
         thread.setStateListener(
             (t, newState, oldState) -> {
@@ -1021,7 +1025,8 @@ public void restore(final Map<TaskId, Task> tasks) {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
 
         final IllegalStateException thrown = assertThrows(
@@ -1060,7 +1065,8 @@ public void shouldShutdownTaskManagerOnCloseWithoutStart() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
         thread.shutdown();
         verify(taskManager);
@@ -1092,7 +1098,8 @@ public void shouldOnlyShutdownOnce() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
         thread.shutdown();
         // Execute the run method. Verification of the mock will check that shutdown was only done once
@@ -1998,7 +2005,8 @@ public void shouldThrowTaskMigratedExceptionHandlingTaskLost() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
 
         consumer.schedulePollTask(() -> {
@@ -2044,7 +2052,8 @@ public void shouldThrowTaskMigratedExceptionHandlingRevocation() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
 
         consumer.schedulePollTask(() -> {
@@ -2096,7 +2105,8 @@ public void shouldCatchHandleCorruptionOnTaskCorruptedExceptionPath() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ) {
             @Override
             void runOnce() {
@@ -2156,7 +2166,8 @@ public void shouldCatchTaskMigratedExceptionOnOnTaskCorruptedExceptionPath() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         ) {
             @Override
             void runOnce() {
@@ -2217,7 +2228,8 @@ public void shouldNotCommitNonRunningNonRestoringTasks() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         );
 
         EasyMock.replay(task1, task2, task3, taskManager);
@@ -2385,7 +2397,8 @@ public void shouldTransmitTaskManagerMetrics() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         );
 
         assertThat(dummyProducerMetrics, is(thread.producerMetrics()));
@@ -2421,7 +2434,8 @@ public void shouldConstructAdminMetrics() {
             new AtomicInteger(),
             new AtomicLong(Long.MAX_VALUE),
             null,
-            HANDLER
+            HANDLER,
+            null
         );
         final MetricName testMetricName = new MetricName("test_metric", "", "", new HashMap<>());
         final Metric testMetric = new KafkaMetric(

streams/src/test/java/org/apache/kafka/streams/state/internals/ThreadCacheTest.java

@@ -471,6 +471,19 @@ public void shouldCalculateSizeInBytes() {
         assertEquals(cache.sizeBytes(), node.size());
     }
 
+    @Test
+    public void shouldResizeAndShrink() {
+        final ThreadCache cache = new ThreadCache(logContext, 10000, new MockStreamsMetrics(new Metrics()));
+        cache.put(namespace, Bytes.wrap(new byte[]{1}), cleanEntry(new byte[]{0}));
+        cache.put(namespace, Bytes.wrap(new byte[]{2}), cleanEntry(new byte[]{0}));
+        cache.put(namespace, Bytes.wrap(new byte[]{3}), cleanEntry(new byte[]{0}));
+        assertEquals(141, cache.sizeBytes());
+        cache.resize(100);
+        assertEquals(94, cache.sizeBytes());
+        cache.put(namespace1, Bytes.wrap(new byte[]{4}), cleanEntry(new byte[]{0}));
+        assertEquals(94, cache.sizeBytes());
+    }
+
     private LRUCacheEntry dirtyEntry(final byte[] key) {
         return new LRUCacheEntry(key, null, true, -1, -1, -1, "");
     }