Refactor MemorySegmentRequest to make concurrent.

Author: leerho

src/main/java/org/apache/datasketches/common/MemorySegmentRequestExample.java

@@ -24,11 +24,12 @@
 import java.lang.foreign.MemorySegment;
 import java.util.Enumeration;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.Objects;
 
 /**
  * This is an example of a possible implementation of the MemorySegmentRequest interface
- * where all requested segments are allocated off-heap.  A local ConcurrentHashMap tracks a newly created confined Arena
- * for every new MemorySegment allocated off-heap. This allows individual segments to be freed
+ * where all requested segments are allocated off-heap.  A local ConcurrentHashMap tracks a newly created
+ * confined Arena for every new MemorySegment allocated off-heap. This allows individual segments to be freed
  * immediately upon receiving the {@link #requestClose(MemorySegment) requestClose(MemorySegment)} call.
  */
 public final class MemorySegmentRequestExample implements MemorySegmentRequest {
@@ -50,24 +51,26 @@ public synchronized MemorySegment request(final long newByteSize) {
 
   @Override
   public synchronized void requestClose(final MemorySegment segKey) {
+    Objects.requireNonNull(segKey, "MemorySegment segKey must not be null");
     final Arena arena = map.get(segKey);
-    if (arena == null) { throw new SketchesArgumentException("Given MemorySegment key is not mapped to an Arena!"); }
-    if (arena.scope().isAlive()) {
-      arena.close();
+    if (arena != null) {
+      if (arena.scope().isAlive()) { arena.close(); }
       map.remove(segKey);
+    } else {
+      //ignore or
+      //throw new SketchesArgumentException("Given MemorySegment key is not mapped to an Arena!");
     }
   }
 
   /**
-   * This cleans up any unclosed, off-heap MemorySegments.
+   * This closes any unclosed, off-heap MemorySegments and removes all mappings from the map.
    */
   public synchronized void cleanup() {
-    for (final Enumeration<Arena> e = map.elements(); e.hasMoreElements();) {
+    for (final Enumeration<Arena> e = map.elements(); e.hasMoreElements(); ) {
       final Arena arena = e.nextElement();
-      if (arena.scope().isAlive()) {
-        arena.close();
-      }
+      if (arena.scope().isAlive()) { arena.close(); }
     }
+    map.clear();
   }
 
 }

src/main/java/org/apache/datasketches/theta/ConcurrentDirectQuickSelectSketch.java

@@ -75,7 +75,9 @@ final class ConcurrentDirectQuickSelectSketch extends DirectQuickSelectSketch
       final double maxConcurrencyError, final MemorySegment dstSeg) {
     super(lgNomLongs, seed, 1.0F, //p
       ResizeFactor.X1, //rf,
-      dstSeg, false); //unionGadget
+      dstSeg,
+      null,
+      false); //unionGadget
 
     volatileThetaLong_ = Long.MAX_VALUE;
     volatileEstimate_ = 0;
@@ -91,6 +93,7 @@ final class ConcurrentDirectQuickSelectSketch extends DirectQuickSelectSketch
     super(sketch.getLgNomLongs(), seed, 1.0F, //p
         ResizeFactor.X1, //rf,
         dstSeg,
+        null,
         false); //unionGadget
 
     exactLimit_ = ConcurrentSharedThetaSketch.computeExactLimit(1L << getLgNomLongs(),
@@ -115,7 +118,7 @@ public double getEstimate() {
 
   @Override
   public boolean isEstimationMode() {
-    return (getRetainedEntries(false) > exactLimit_) || super.isEstimationMode();
+    return getRetainedEntries(false) > exactLimit_ || super.isEstimationMode();
   }
 
   @Override
@@ -164,7 +167,7 @@ public long getExactLimit() {
   @Override
   public boolean startEagerPropagation() {
     while (!sharedPropagationInProgress_.compareAndSet(false, true)) { /* busy wait till free */ }
-    return (!isEstimationMode());// no eager propagation is allowed in estimation mode
+    return !isEstimationMode();// no eager propagation is allowed in estimation mode
   }
 
   @Override
@@ -206,8 +209,8 @@ public void initBgPropagationService() {
   public boolean propagate(final AtomicBoolean localPropagationInProgress,
                            final Sketch sketchIn, final long singleHash) {
     final long epoch = epoch_;
-    if ((singleHash != NOT_SINGLE_HASH)                   // namely, is a single hash and
-        && (getRetainedEntries(false) < exactLimit_)) {   // a small sketch then propagate myself (blocking)
+    if (singleHash != NOT_SINGLE_HASH                   // namely, is a single hash and
+        && getRetainedEntries(false) < exactLimit_) {   // a small sketch then propagate myself (blocking)
       if (!startEagerPropagation()) {
         endPropagation(localPropagationInProgress, true);
         return false;

src/main/java/org/apache/datasketches/theta/DirectQuickSelectSketch.java

@@ -60,6 +60,7 @@
 import java.lang.foreign.MemorySegment;
 
 import org.apache.datasketches.common.Family;
+import org.apache.datasketches.common.MemorySegmentRequest;
 import org.apache.datasketches.common.ResizeFactor;
 import org.apache.datasketches.common.SketchesArgumentException;
 import org.apache.datasketches.common.SuppressFBWarnings;
@@ -81,15 +82,19 @@
 class DirectQuickSelectSketch extends DirectQuickSelectSketchR {
   private static final double DQS_RESIZE_THRESHOLD  = 15.0 / 16.0; //tuned for space
   int hashTableThreshold_; //computed and mutable, kept only on heap, never serialized.
+  private final MemorySegmentRequest mSegReq;
 
   /**
-   * Construct this sketch as a result of a wrap operation where the given MemorySegment already has a sketch image.
-   * @param wseg the given MemorySegment that has a sketch image.
+   * Construct this sketch as a result of a wrap operation where the given MemorySegment already has an updatable sketch image.
+   * @param wseg the given MemorySegment that has an updatable sketch image.
+   * @param mSegReq an implementation of the MemorySegmentRequest interface or null.
    * @param seed <a href="{@docRoot}/resources/dictionary.html#seed">See Update Hash Seed</a>.
    */
   private DirectQuickSelectSketch(
       final MemorySegment wseg,
+      final MemorySegmentRequest mSegReq,
       final long seed) {
+    this.mSegReq = mSegReq == null ? MemorySegmentRequest.DEFAULT : mSegReq;
     super(wseg, seed);
   }
 
@@ -105,6 +110,7 @@ private DirectQuickSelectSketch(
    * <a href="{@docRoot}/resources/dictionary.html#resizeFactor">See Resize Factor</a>
    * @param dstSeg the given MemorySegment object destination. It cannot be null.
    * It will be cleared prior to use.
+   * @param mSegReq an implementation of the MemorySegmentRequest interface or null.
    * @param unionGadget true if this sketch is implementing the Union gadget function.
    * Otherwise, it is behaving as a normal QuickSelectSketch.
    */
@@ -114,6 +120,7 @@ private DirectQuickSelectSketch(
       final float p,
       final ResizeFactor rf,
       final MemorySegment dstSeg,
+      final MemorySegmentRequest mSegReq,
       final boolean unionGadget) {
 
     //Choose family, preambleLongs
@@ -129,7 +136,7 @@ private DirectQuickSelectSketch(
     final long curSegCapBytes = dstSeg.byteSize();
     if (curSegCapBytes < minReqBytes) {
       throw new SketchesArgumentException(
-        "MemorySegment capacity is too small: " + curSegCapBytes + " < " + minReqBytes);
+        "MemorySegment capacity is less than minimum required: " + curSegCapBytes + " < " + minReqBytes);
     }
 
     //@formatter:off
@@ -153,17 +160,22 @@ private DirectQuickSelectSketch(
     //clear hash table area
     dstSeg.asSlice(preambleLongs << 3, Long.BYTES << lgArrLongs).fill((byte)0);
     hashTableThreshold_ = getOffHeapHashTableThreshold(lgNomLongs, lgArrLongs);
+    this.mSegReq = mSegReq == null ? MemorySegmentRequest.DEFAULT : mSegReq;
     super(dstSeg, seed);
   }
 
   /**
-   * Wrap a sketch around the given source MemorySegment containing sketch data that originated from
-   * this sketch.
+   * Wrap a sketch around the given source MemorySegment containing sketch data that originated from this sketch.
    * @param srcSeg The given MemorySegment object must be in hash table form and not read only.
+   * @param mSegReq an implementation of the MemorySegmentRequest interface or null.
    * @param seed <a href="{@docRoot}/resources/dictionary.html#seed">See Update Hash Seed</a>
    * @return instance of this sketch
    */
-  static DirectQuickSelectSketch writableWrap(final MemorySegment srcSeg, final long seed) {
+  //called from UnionImpl and UpdateSketch
+  static DirectQuickSelectSketch writableWrap(
+      final MemorySegment srcSeg,
+      final MemorySegmentRequest mSegReq,
+      final long seed) {
     final int preambleLongs = extractPreLongs(srcSeg);                  //byte 0
     final int lgNomLongs = extractLgNomLongs(srcSeg);                   //byte 3
     final int lgArrLongs = extractLgArrLongs(srcSeg);                   //byte 4
@@ -176,7 +188,7 @@ static DirectQuickSelectSketch writableWrap(final MemorySegment srcSeg, final lo
       insertLgResizeFactor(srcSeg, ResizeFactor.X2.lg());
     }
 
-    final DirectQuickSelectSketch dqss = new DirectQuickSelectSketch(srcSeg, seed);
+    final DirectQuickSelectSketch dqss = new DirectQuickSelectSketch(srcSeg, mSegReq, seed);
     dqss.hashTableThreshold_ = getOffHeapHashTableThreshold(lgNomLongs, lgArrLongs);
     return dqss;
   }
@@ -185,14 +197,19 @@ static DirectQuickSelectSketch writableWrap(final MemorySegment srcSeg, final lo
    * Fast-wrap a sketch around the given source MemorySegment containing sketch data that originated from
    * this sketch.  This does NO validity checking of the given MemorySegment.
    * @param srcSeg The given MemorySegment must be in hash table form and not read only.
+   * @param mSegReq an implementation of the MemorySegmentRequest interface or null.
    * @param seed <a href="{@docRoot}/resources/dictionary.html#seed">See Update Hash Seed</a>
    * @return instance of this sketch
    */
-  static DirectQuickSelectSketch fastWritableWrap(final MemorySegment srcSeg, final long seed) {
+  //called from UnionImpl <- Union
+  static DirectQuickSelectSketch fastWritableWrap(
+      final MemorySegment srcSeg,
+      final MemorySegmentRequest mSegReq,
+      final long seed) {
     final int lgNomLongs = extractLgNomLongs(srcSeg);                   //byte 3
     final int lgArrLongs = extractLgArrLongs(srcSeg);                   //byte 4
 
-    final DirectQuickSelectSketch dqss = new DirectQuickSelectSketch(srcSeg, seed);
+    final DirectQuickSelectSketch dqss = new DirectQuickSelectSketch(srcSeg, mSegReq, seed);
     dqss.hashTableThreshold_ = getOffHeapHashTableThreshold(lgNomLongs, lgArrLongs);
     return dqss;
   }
@@ -205,7 +222,7 @@ static DirectQuickSelectSketch fastWritableWrap(final MemorySegment srcSeg, fina
   public UpdateSketch rebuild() {
     final int lgNomLongs = getLgNomLongs();
     final int preambleLongs = wseg_.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) & 0X3F;
-    if (getRetainedEntries(true) > (1 << lgNomLongs)) {
+    if (getRetainedEntries(true) > 1 << lgNomLongs) {
       quickSelectAndRebuild(wseg_, preambleLongs, lgNomLongs);
     }
     return this;
@@ -279,10 +296,13 @@ UpdateReturnState hashUpdate(final long hash) {
           tgtLgArrLongs = Math.min(lgArrLongs + lgRF, lgNomLongs + 1);
           final int tgtArrBytes = 8 << tgtLgArrLongs;
           final int reqBytes = tgtArrBytes + preBytes;
-          final MemorySegment newDstSeg = MemorySegment.ofArray(new byte[reqBytes]); //always on-heap //TODO ADD MemSegReq
+
+          final MemorySegment newDstSeg = mSegReq.request(reqBytes);
 
           moveAndResize(wseg_, preambleLongs, lgArrLongs, newDstSeg, tgtLgArrLongs, thetaLong);
+          final MemorySegment oldSeg = wseg_;
           wseg_ = newDstSeg;
+          mSegReq.requestClose(oldSeg);
 
           hashTableThreshold_ = getOffHeapHashTableThreshold(lgNomLongs, tgtLgArrLongs);
           return InsertedCountIncrementedResized;

src/main/java/org/apache/datasketches/theta/PreambleUtil.java

@@ -272,7 +272,7 @@ static String preambleToString(final MemorySegment seg) {
 
     //Flags
     final int flags = extractFlags(seg);
-    final String flagsStr = (flags) + ", 0x" + (Integer.toHexString(flags)) + ", "
+    final String flagsStr = flags + ", 0x" + Integer.toHexString(flags) + ", "
         + zeroPad(Integer.toBinaryString(flags), 8);
     final boolean readOnly = (flags & READ_ONLY_FLAG_MASK) > 0;
     final boolean empty = (flags & EMPTY_FLAG_MASK) > 0;
@@ -377,11 +377,11 @@ else if (preLongs == 3) {
   //@formatter:on
 
   static int extractPreLongs(final MemorySegment seg) {
-    return seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) & 0X3F;
+    return seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) & 0X3F; //for SerVer 1,2,3
   }
 
   static int extractLgResizeFactor(final MemorySegment seg) {
-    return (seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) >>> LG_RESIZE_FACTOR_BIT) & 0X3;
+    return seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) >>> LG_RESIZE_FACTOR_BIT & 0X3;
   }
 
   static int extractLgResizeRatioV1(final MemorySegment seg) {
@@ -463,7 +463,7 @@ static void insertLgResizeFactor(final MemorySegment seg, final int rf) {
     final int curByte = seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE) & 0xFF;
     final int shift = LG_RESIZE_FACTOR_BIT; // shift in bits
     final int mask = 3;
-    final byte newByte = (byte) (((rf & mask) << shift) | (~(mask << shift) & curByte));
+    final byte newByte = (byte) ((rf & mask) << shift | ~(mask << shift) & curByte);
     seg.set(JAVA_BYTE, PREAMBLE_LONGS_BYTE, newByte);
   }
 
@@ -520,7 +520,7 @@ static void clearEmpty(final MemorySegment seg) {
   }
 
   static boolean isEmptyFlag(final MemorySegment seg) {
-    return ((extractFlags(seg) & EMPTY_FLAG_MASK) > 0);
+    return (extractFlags(seg) & EMPTY_FLAG_MASK) > 0;
   }
 
   /**

src/main/java/org/apache/datasketches/theta/Sketch.java

@@ -340,7 +340,7 @@ public static int getMaxUpdateSketchBytes(final int nomEntries) {
    * @return the maximum number of storage bytes required for a UpdateSketch with the given lgNomEntries
    */
   public static int getUpdateSketchMaxBytes(final int lgNomEntries) {
-    return (1 << lgNomEntries << 4) + (Family.QUICKSELECT.getMaxPreLongs() << 3);
+    return (16 << lgNomEntries) + (Family.QUICKSELECT.getMaxPreLongs() << 3);
   }
 
   /**

src/main/java/org/apache/datasketches/theta/Sketches.java

@@ -355,12 +355,12 @@ public static UpdateSketch wrapUpdateSketch(final MemorySegment srcSeg) {
    * @return {@link UpdateSketch UpdateSketch}
    */
   public static UpdateSketch wrapUpdateSketch(final MemorySegment srcSeg, final long expectedSeed) {
-    return UpdateSketch.wrap(srcSeg, expectedSeed);
+    return UpdateSketch.wrap(srcSeg, null, expectedSeed);
   }
 
   //Restricted static methods
 
-  static void checkIfValidThetaSketch(final MemorySegment srcSeg) {
+  private static void checkIfValidThetaSketch(final MemorySegment srcSeg) {
     final int fam = srcSeg.get(JAVA_BYTE, FAMILY_BYTE);
     if (!Sketch.isValidSketchID(fam)) {
      throw new SketchesArgumentException("Source MemorySegment not a valid Sketch. Family: "
@@ -371,7 +371,7 @@ static void checkIfValidThetaSketch(final MemorySegment srcSeg) {
   static boolean getEmpty(final MemorySegment srcSeg) {
     final int serVer = srcSeg.get(JAVA_BYTE, SER_VER_BYTE);
     if (serVer == 1) {
-      return ((getThetaLong(srcSeg) == Long.MAX_VALUE) && (getRetainedEntries(srcSeg) == 0));
+      return getThetaLong(srcSeg) == Long.MAX_VALUE && getRetainedEntries(srcSeg) == 0;
     }
     return (srcSeg.get(JAVA_BYTE, FLAGS_BYTE) & EMPTY_FLAG_MASK) != 0; //for SerVer 2 & 3
   }
@@ -384,7 +384,7 @@ static int getRetainedEntries(final MemorySegment srcSeg) {
     final int serVer = srcSeg.get(JAVA_BYTE, SER_VER_BYTE);
     if (serVer == 1) {
       final int entries = srcSeg.get(JAVA_INT_UNALIGNED, RETAINED_ENTRIES_INT);
-      if ((getThetaLong(srcSeg) == Long.MAX_VALUE) && (entries == 0)) {
+      if (getThetaLong(srcSeg) == Long.MAX_VALUE && entries == 0) {
         return 0;
       }
       return entries;
@@ -401,6 +401,6 @@ static int getRetainedEntries(final MemorySegment srcSeg) {
 
   static long getThetaLong(final MemorySegment srcSeg) {
     final int preLongs = getPreambleLongs(srcSeg);
-    return (preLongs < 3) ? Long.MAX_VALUE : srcSeg.get(JAVA_LONG_UNALIGNED, THETA_LONG); //for SerVer 1,2,3
+    return preLongs < 3 ? Long.MAX_VALUE : srcSeg.get(JAVA_LONG_UNALIGNED, THETA_LONG); //for SerVer 1,2,3
   }
 }

src/main/java/org/apache/datasketches/theta/UnionImpl.java

@@ -105,7 +105,7 @@ static UnionImpl initNewDirectInstance(
       final ResizeFactor rf,
       final MemorySegment dstSeg) {
     final UpdateSketch gadget = //create with UNION family
-        new DirectQuickSelectSketch(lgNomLongs, seed, p, rf, dstSeg, true);
+        new DirectQuickSelectSketch(lgNomLongs, seed, p, rf, dstSeg, null, true);
     final UnionImpl unionImpl = new UnionImpl(gadget, seed);
     unionImpl.unionThetaLong_ = gadget.getThetaLong();
     unionImpl.unionEmpty_ = gadget.isEmpty();
@@ -142,7 +142,7 @@ static UnionImpl fastWrapInstance(final MemorySegment srcSeg, final long expecte
     Family.UNION.checkFamilyID(extractFamilyID(srcSeg));
     final UpdateSketch gadget = srcSeg.isReadOnly()
         ? DirectQuickSelectSketchR.fastReadOnlyWrap(srcSeg, expectedSeed)
-        : DirectQuickSelectSketch.fastWritableWrap(srcSeg, expectedSeed);
+        : DirectQuickSelectSketch.fastWritableWrap(srcSeg, null, expectedSeed);
     final UnionImpl unionImpl = new UnionImpl(gadget, expectedSeed);
     unionImpl.unionThetaLong_ = extractUnionThetaLong(srcSeg);
     unionImpl.unionEmpty_ = PreambleUtil.isEmptyFlag(srcSeg);
@@ -151,17 +151,17 @@ static UnionImpl fastWrapInstance(final MemorySegment srcSeg, final long expecte
 
   /**
    * Wrap a Union object around a Union MemorySegment object containing data.
-   * Called by SetOperation.
    * @param srcSeg The source MemorySegment object.
    * @param expectedSeed the seed used to validate the given MemorySegment image.
    * <a href="{@docRoot}/resources/dictionary.html#seed">See seed</a>
    * @return this class
    */
+  //Called by SetOperation and Union
   static UnionImpl wrapInstance(final MemorySegment srcSeg, final long expectedSeed) {
     Family.UNION.checkFamilyID(extractFamilyID(srcSeg));
     final UpdateSketch gadget = srcSeg.isReadOnly()
         ? DirectQuickSelectSketchR.readOnlyWrap(srcSeg, expectedSeed)
-        : DirectQuickSelectSketch.writableWrap(srcSeg, expectedSeed);
+        : DirectQuickSelectSketch.writableWrap(srcSeg, null, expectedSeed);
     final UnionImpl unionImpl = new UnionImpl(gadget, expectedSeed);
     unionImpl.unionThetaLong_ = extractUnionThetaLong(srcSeg);
     unionImpl.unionEmpty_ = PreambleUtil.isEmptyFlag(srcSeg);
@@ -269,7 +269,7 @@ public CompactSketch union(final Sketch sketchA, final Sketch sketchB, final boo
   public void union(final Sketch sketchIn) {
     //UNION Empty Rule: AND the empty states.
 
-    if ((sketchIn == null) || sketchIn.isEmpty()) {
+    if (sketchIn == null || sketchIn.isEmpty()) {
       //null and empty is interpreted as (Theta = 1.0, count = 0, empty = T).  Nothing changes
       return;
     }
@@ -287,7 +287,7 @@ public void union(final Sketch sketchIn) {
     final HashIterator it = sketchIn.iterator();
     while (it.next()) {
       final long hash = it.get();
-      if ((hash < unionThetaLong_) && (hash < gadget_.getThetaLong())) {
+      if (hash < unionThetaLong_ && hash < gadget_.getThetaLong()) {
         gadget_.hashUpdate(hash); // backdoor update, hash function is bypassed
       } else if (isOrdered) { break; }
     }