ISPN-7922 OffHeap stream causes too many entries in memory at once

* Make sure that we only have 1 flatMap * Reduce invocation count for lots of writes
infinispan · Jun 15, 2017 · 3b33eda · 3b33eda
1 parent 0f5273f
commit 3b33eda
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 33 deletions.
diff --git a/core/src/main/java/org/infinispan/container/offheap/OffHeapDataContainer.java b/core/src/main/java/org/infinispan/container/offheap/OffHeapDataContainer.java
@@ -11,8 +11,8 @@
 import java.util.concurrent.locks.Lock;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.stream.IntStream;
-import java.util.stream.LongStream;
 import java.util.stream.Stream;
 
 import org.infinispan.commons.marshall.WrappedByteArray;
@@ -562,37 +562,27 @@ public void executeTask(KeyValueFilter<? super WrappedBytes, ? super WrappedByte
    }
 
    private Stream<InternalCacheEntry<WrappedBytes, WrappedBytes>> entryStream() {
-      int limit = memoryAddressCount / lockCount;
-      return IntStream.range(0, lockCount)
-            // REALLY REALLY stupid there is no flatMapToObj on IntStream...
-            .boxed()
-            .flatMap(l -> {
-               int value = l;
-               return LongStream.iterate(value, i -> i + lockCount).limit(limit)
-                     .boxed()
-                     .flatMap(a -> {
-                     Lock lock = locks.getLockWithOffset(value).readLock();
-                     lock.lock();
-                     try {
-                        checkDeallocation();
-                        long address = memoryLookup.getMemoryAddressOffset(a.intValue());
-                        if (address == 0) {
-                           return Stream.empty();
-                        }
-                        Stream.Builder<InternalCacheEntry<WrappedBytes, WrappedBytes>> builder = Stream.builder();
-                        while (address != 0) {
-                           long nextAddress;
-                           do {
-                              nextAddress = offHeapEntryFactory.getNextLinkedPointerAddress(address);
-                              builder.accept(offHeapEntryFactory.fromMemory(address));
-                           } while ((address = nextAddress) != 0);
-                        }
-                        return builder.build();
-                     } finally {
-                        lock.unlock();
-                     }
-                  });
-            });
+      return IntStream.range(0, memoryAddressCount)
+            .mapToObj(a -> {
+               Lock lock = locks.getLockWithOffset(a % lockCount).readLock();
+               lock.lock();
+               try {
+                  checkDeallocation();
+                  long address = memoryLookup.getMemoryAddressOffset(a);
+                  if (address == 0) {
+                     return null;
+                  }
+                  Stream.Builder<InternalCacheEntry<WrappedBytes, WrappedBytes>> builder = Stream.builder();
+                  long nextAddress;
+                  do {
+                     nextAddress = offHeapEntryFactory.getNextLinkedPointerAddress(address);
+                     builder.accept(offHeapEntryFactory.fromMemory(address));
+                  } while ((address = nextAddress) != 0);
+                  return builder.build();
+               } finally {
+                  lock.unlock();
+               }
+            }).flatMap(Function.identity());
    }
 
    @Override

diff --git a/core/src/test/java/org/infinispan/container/offheap/OffHeapSingleNodeTest.java b/core/src/test/java/org/infinispan/container/offheap/OffHeapSingleNodeTest.java
@@ -84,7 +84,7 @@ public void testLockOnExecuteTask() throws InterruptedException, TimeoutExceptio
    public void testLotsOfWrites() {
       Cache<String, String> cache = cache(0);
 
-      for (int i = 0; i < 100000; ++i) {
+      for (int i = 0; i < 5_000; ++i) {
          cache.put("key" + i, "value" + i);
       }