Updated MemorySegmentRequestExample to use ConcurrentHashMap.

Author: leerho

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

@@ -29,15 +29,13 @@
 public interface MemorySegmentRequest {
 
   /**
-   * Request a new MemorySegment with the given <i>newByteSize</i>.
+   * Request a new heap MemorySegment with the given <i>newByteSize</i>.
    * Because we do not have a reference to an Arena, the default here is to
    * allocate a new MemorySegment on the heap.  It is up to the user to override this as appropriate.
-   * @param prevSeg the previous MemorySegment to be possibly closed here or by using the separate
-   * {@link #requestClose requestClose} method. This is included for convenience, it may be null.
    * @param newByteSize The new <i>byteSize</i> being requested.
    * @return new MemorySegment with the requested <i>byteSize</i>.
    */
-  default MemorySegment request(final MemorySegment prevSeg, final long newByteSize) {
+  default MemorySegment request(final long newByteSize) {
     if (newByteSize > Integer.MAX_VALUE) {
       throw new SketchesArgumentException("Requested size in bytes exceeds Integer.MAX_VALUE.");
     }

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

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.datasketches.common;
+
+import java.lang.foreign.Arena;
+
+import java.lang.foreign.MemorySegment;
+import java.util.Enumeration;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * 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
+ * immediately upon receiving the {@link #requestClose(MemorySegment) requestClose(MemorySegment)} call.
+ */
+public final class MemorySegmentRequestExample implements MemorySegmentRequest {
+  private final ConcurrentHashMap<MemorySegment, Arena> map = new ConcurrentHashMap<>();
+
+  /**
+   * Request a new off-heap MemorySegment with the given <i>newByteSeze</i>.
+   * An internal confined Arena is created to exclusively manage the new segment and it is associated
+   * with the new segment with a ConcurrentHashMap.
+   */
+  @Override
+  public synchronized MemorySegment request(final long newByteSize) {
+    final Arena arena = Arena.ofConfined();
+    final MemorySegment seg = arena.allocate(newByteSize);
+    map.put(seg, arena);
+    return seg;
+
+  }
+
+  @Override
+  public synchronized void requestClose(final MemorySegment segKey) {
+    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();
+      map.remove(segKey);
+    }
+  }
+
+  /**
+   * This cleans up any unclosed, off-heap MemorySegments.
+   */
+  public synchronized void cleanup() {
+    for (final Enumeration<Arena> e = map.elements(); e.hasMoreElements();) {
+      final Arena arena = e.nextElement();
+      if (arena.scope().isAlive()) {
+        arena.close();
+      }
+    }
+  }
+
+}

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

@@ -357,7 +357,7 @@ static MemorySegment memorySegmentSpaceMgmt(
       if (mSegReq == null) {
         mSegReq = MemorySegmentRequest.DEFAULT;
       }
-      final MemorySegment newSeg = mSegReq.request(oldWseg, requiredSketchBytes);
+      final MemorySegment newSeg = mSegReq.request(requiredSketchBytes);
       MemorySegment.copy(oldWseg, 0, newSeg, 0, DATA_START_ADR); //copy preamble (first 20 bytes)
       mSegReq.requestClose(oldWseg);
       return newSeg;

src/main/java/org/apache/datasketches/kll/KllHelper.java

@@ -358,7 +358,7 @@ private MemorySegment growCombinedSegBuffer(final int itemSpaceNeeded) {
 
     mSegReq_ = (mSegReq_ == null) ? MemorySegmentRequest.DEFAULT : mSegReq_;
 
-    final MemorySegment newSeg = mSegReq_.request(seg_, needBytes);
+    final MemorySegment newSeg = mSegReq_.request(needBytes);
     MemorySegment.copy(seg_, 0, newSeg, 0, segBytes);
     mSegReq_.requestClose(seg_);
     return newSeg;

src/main/java/org/apache/datasketches/quantiles/DirectUpdateDoublesSketch.java

@@ -27,7 +27,7 @@
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 
-import org.apache.datasketches.common.MemorySegmentRequestExtension;
+import org.apache.datasketches.common.MemorySegmentRequestExample;
 import org.testng.annotations.Test;
 
 
@@ -49,7 +49,7 @@ public void checkMemorySegmentRequestExtension() {
     final MemorySegment seg = arena.allocate(numBytes);
 
     //Use the custom extension of the MemorySegmentRequest interface.
-    final MemorySegmentRequestExtension mSegReqExt = new MemorySegmentRequestExtension();
+    final MemorySegmentRequestExample mSegReqExt = new MemorySegmentRequestExample();
 
     //Create a new KllLongsSketch and pass the custom extension
     final KllLongsSketch sk = KllLongsSketch.newDirectInstance(k, seg, mSegReqExt);

src/test/java/org/apache/datasketches/kll/KllMemorySegmentRequestApp.java

@@ -25,7 +25,7 @@
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 
-import org.apache.datasketches.common.MemorySegmentRequestExtension;
+import org.apache.datasketches.common.MemorySegmentRequestExample;
 import org.testng.annotations.Test;
 
 public class ClassicQuantilesMemorySegmentRequestApp {
@@ -47,7 +47,7 @@ public void checkMemorySegmentRequestExtension() {
     final MemorySegment seg = arena.allocate(initalBytes);
 
     //Use the custom extension of the MemorySegmentRequest interface.
-    final MemorySegmentRequestExtension mSegReqExt = new MemorySegmentRequestExtension();
+    final MemorySegmentRequestExample mSegReqExt = new MemorySegmentRequestExample();
 
     //Create a new KllLongsSketch and pass the custom extension
     final DoublesSketchBuilder bldr = DoublesSketch.builder().setK(k);