Fix RocksDB native memory leak in stateful applications (#32)

* Fix RocksDB native memory leak

* Fix unit tests

Author: laxmanchekka

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/rocksdb/RocksDBCacheProvider.java

@@ -13,6 +13,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import java.util.Map;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.streams.state.internals.BlockBasedTableConfigWithAccessibleCache;
 import org.rocksdb.BlockBasedTableConfig;
 import org.rocksdb.Cache;
 import org.rocksdb.LRUCache;
@@ -107,7 +108,7 @@ protected synchronized void initCache(Options options, Map<String, Object> confi
           cacheTotalCapacity / (1024 * 1024), writeBuffersRatio, highPriorityPoolRatio);
     }
 
-    final BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) options.tableFormatConfig();
+    final BlockBasedTableConfigWithAccessibleCache tableConfig = (BlockBasedTableConfigWithAccessibleCache) options.tableFormatConfig();
 
     // ######### Block cache (Read buffers) #########
     if (configs.containsKey(BLOCK_SIZE)) {
@@ -133,7 +134,12 @@ protected synchronized void initCache(Options options, Map<String, Object> confi
 
     options.setWriteBufferManager(writeBufferManager);
 
-    tableConfig.setBlockCache(cache);
+    final Cache oldCache = tableConfig.blockCache();
+    if(oldCache != null) {
+      LOG.info("Releasing old rocksdb cache before setting shared cache.");
+      oldCache.close();
+    }
+    tableConfig.setBlockCache(this.cache);
     options.setTableFormatConfig(tableConfig);
   }
 

kafka-streams-framework/src/test/java/org/hypertrace/core/kafkastreams/framework/rocksdb/BoundedMemoryConfigSetterTest.java

@@ -17,6 +17,7 @@
 import java.util.stream.Stream;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.streams.state.RocksDBConfigSetter;
+import org.apache.kafka.streams.state.internals.BlockBasedTableConfigWithAccessibleCache;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -26,6 +27,7 @@
 import org.rocksdb.CompactionStyle;
 import org.rocksdb.CompressionType;
 import org.rocksdb.InfoLogLevel;
+import org.rocksdb.LRUCache;
 import org.rocksdb.Options;
 
 class BoundedMemoryConfigSetterTest {
@@ -41,7 +43,9 @@ public void setUp() {
     RocksDBCacheProvider.get().testDestroy();
     options = new Options();
     configs = new HashMap<>();
-    tableConfig = new BlockBasedTableConfig();
+    tableConfig = new BlockBasedTableConfigWithAccessibleCache();
+    // mimic kafka streams
+    tableConfig.setBlockCache(new LRUCache(32 * 1024));
     options.setTableFormatConfig(tableConfig);
     configSetter = new BoundedMemoryConfigSetter();
   }