[ci] use netty bufferPools for direct buffers (#3437)

Signed-off-by: Itai Segall <itai.segall@digitalasset.com>

Author: isegall-da

apps/scan/src/main/scala/org/lfdecentralizedtrust/splice/scan/store/bulk/ZstdGroupedWeight.scala

@@ -4,11 +4,11 @@
 package org.lfdecentralizedtrust.splice.scan.store.bulk
 
 import com.github.luben.zstd.ZstdDirectBufferCompressingStreamNoFinalizer
+import io.netty.buffer.PooledByteBufAllocator
 import org.apache.pekko.stream.{Attributes, FlowShape, Inlet, Outlet}
 import org.apache.pekko.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
 import org.apache.pekko.util.ByteString
 
-import java.nio.ByteBuffer
 import java.util.concurrent.atomic.AtomicReference
 
 /** A Pekko GraphStage that zstd-compresses a stream of bytestrings, and splits the output into zstd objects of size (minWeight + delta).
@@ -41,45 +41,44 @@ case class ZstdGroupedWeight(minSize: Long) extends GraphStage[FlowShape[ByteStr
   }
 
   class ZSTD(
-      val tmpBuffer: ByteBuffer,
-      val compressionLevel: Int = 3,
+      val compressionLevel: Int = 3
   ) extends AutoCloseable {
 
+    val bufferAllocator = PooledByteBufAllocator.DEFAULT
+    val tmpBuffer = bufferAllocator.directBuffer(zstdTmpBufferSize)
+    val tmpNioBuffer = tmpBuffer.nioBuffer(0, tmpBuffer.capacity())
     val compressingStream =
-      new ZstdDirectBufferCompressingStreamNoFinalizer(tmpBuffer, compressionLevel)
+      new ZstdDirectBufferCompressingStreamNoFinalizer(tmpNioBuffer, compressionLevel)
 
     def compress(input: ByteString): ByteString = {
-      // TODO(#3429): use a buffer pool to avoid allocating a new ByteBuffer for each compress call
-      val inputBB = ByteBuffer.allocateDirect(input.size)
-      inputBB.put(input.toArrayUnsafe())
-      inputBB.flip()
-      compressingStream.compress(inputBB)
+      val inputBB = bufferAllocator.directBuffer(input.size)
+      inputBB.writeBytes(input.toArrayUnsafe())
+      compressingStream.compress(inputBB.nioBuffer())
+      inputBB.release()
       compressingStream.flush()
-      tmpBuffer.flip()
-      val result = ByteString.fromByteBuffer(tmpBuffer)
-      tmpBuffer.clear()
+      tmpNioBuffer.flip()
+      val result = ByteString.fromByteBuffer(tmpNioBuffer)
+      tmpNioBuffer.clear()
       result
     }
 
     def zstdFinish(): ByteString = {
-      tmpBuffer.flip()
-      val result = ByteString.fromByteBuffer(tmpBuffer)
-      tmpBuffer.clear()
+      tmpNioBuffer.flip()
+      val result = ByteString.fromByteBuffer(tmpNioBuffer)
+      tmpNioBuffer.clear()
       compressingStream.close()
       result
     }
 
     override def close(): Unit = {
       compressingStream.close()
+      val _ = tmpBuffer.release()
     }
   }
 
   override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
     new GraphStageLogic(shape) with InHandler with OutHandler {
-      // TODO(#3429): consider implementing a pool of tmp buffers to avoid allocating a new one for each stage,
-      //   and moving some initialization into preStart(), otherwise we allocate even if the stream never runs or fails before starting.
-      private val tmpBuffer = ByteBuffer.allocateDirect(zstdTmpBufferSize)
-      private val zstd = new AtomicReference[ZSTD](new ZSTD(tmpBuffer, 3))
+      private val zstd = new AtomicReference[ZSTD](new ZSTD(3))
       private val state: AtomicReference[State] = new AtomicReference[State](State.empty())
 
       override def postStop(): Unit = {
@@ -90,9 +89,8 @@ case class ZstdGroupedWeight(minSize: Long) extends GraphStage[FlowShape[ByteStr
       }
 
       private def reset(): Unit = {
-        tmpBuffer.clear()
         zstd.get().close()
-        zstd.set(new ZSTD(tmpBuffer, 3))
+        zstd.set(new ZSTD(3))
         state.set(State.empty())
       }
 

apps/scan/src/test/scala/org/lfdecentralizedtrust/splice/scan/store/AcsSnapshotBulkStorageTest.scala

@@ -11,6 +11,7 @@ import com.digitalasset.canton.tracing.TraceContext
 import com.digitalasset.canton.{HasActorSystem, HasExecutionContext}
 import com.github.luben.zstd.ZstdDirectBufferDecompressingStream
 import com.google.protobuf.ByteString
+import io.netty.buffer.PooledByteBufAllocator
 import org.lfdecentralizedtrust.splice.scan.admin.http.CompactJsonScanHttpEncodings
 import org.lfdecentralizedtrust.splice.scan.store.AcsSnapshotStore.QueryAcsSnapshotResult
 import org.lfdecentralizedtrust.splice.scan.store.bulk.{
@@ -19,15 +20,14 @@ import org.lfdecentralizedtrust.splice.scan.store.bulk.{
   S3BucketConnection,
   S3Config,
 }
-import org.lfdecentralizedtrust.splice.store.{Limit, StoreTest, HardLimit}
+import org.lfdecentralizedtrust.splice.store.{HardLimit, Limit, StoreTest}
 import org.lfdecentralizedtrust.splice.store.events.SpliceCreatedEvent
 import org.lfdecentralizedtrust.splice.util.{EventId, PackageQualifiedName, ValueJsonCodecCodegen}
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.s3.model.{ListObjectsRequest, S3Object}
 
 import java.net.URI
-import java.nio.ByteBuffer
 import java.nio.charset.StandardCharsets
 import scala.concurrent.Future
 import scala.jdk.FutureConverters.*
@@ -113,15 +113,20 @@ class AcsSnapshotBulkStorageTest extends StoreTest with HasExecutionContext with
   def readUncompressAndDecode(
       s3BucketConnection: S3BucketConnection
   )(s3obj: S3Object): Array[httpApi.CreatedEvent] = {
+    val bufferAllocator = PooledByteBufAllocator.DEFAULT
     val compressed = s3BucketConnection.readFullObject(s3obj.key()).futureValue
-    val compressedDirect = ByteBuffer.allocateDirect(compressed.capacity())
-    val uncompressed = ByteBuffer.allocateDirect(compressed.capacity() * 200)
-    compressedDirect.put(compressed)
-    compressedDirect.flip()
-    Using(new ZstdDirectBufferDecompressingStream(compressedDirect)) { _.read(uncompressed) }
-    uncompressed.flip()
-    val allContractsStr = StandardCharsets.UTF_8.newDecoder().decode(uncompressed).toString
+    val compressedDirect = bufferAllocator.directBuffer(compressed.capacity())
+    val uncompressedDirect = bufferAllocator.directBuffer(compressed.capacity() * 200)
+    val uncompressedNio = uncompressedDirect.nioBuffer(0, uncompressedDirect.capacity())
+    compressedDirect.writeBytes(compressed)
+    Using(new ZstdDirectBufferDecompressingStream(compressedDirect.nioBuffer())) {
+      _.read(uncompressedNio)
+    }
+    uncompressedNio.flip()
+    val allContractsStr = StandardCharsets.UTF_8.newDecoder().decode(uncompressedNio).toString
     val allContracts = allContractsStr.split("\n")
+    compressedDirect.release()
+    uncompressedDirect.release()
     allContracts.map(io.circe.parser.decode[httpApi.CreatedEvent](_).value)
   }