repo-sync-2025-08-25T14:15:37+0800

build/Dockerfiles/dataproxy.Dockerfile

@@ -23,7 +23,13 @@ RUN ln -s ${JAVA_HOME}/lib/libjli.so /lib64
 COPY dataproxy-server/target/dataproxy-server-0.0.1-SNAPSHOT.jar dataproxy.jar
 COPY libs/*.jar libs/
 
-ENV JAVA_OPTS="-server -XX:+UseG1GC -XX:+DisableExplicitGC -XX:InitiatingHeapOccupancyPercent=68 -Xlog:gc*=info:file=gc.log:time,tags:filecount=5,filesize=10M -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/app/log -XX:ErrorFile=/app/log/hs_err_pid%p.log"
+# At present, when the concurrency of the task is not high,
+# the heap memory consumption is not large, the maximum setting is 512M,
+# if the resources are sufficient, and the concurrency can be appropriately increased,
+# but the MaxDirectMemorySize needs to be increased according to the situation,
+# the current default setting is 1536m,
+# You can override this configuration by adding environment variables
+ENV JAVA_OPTS="-server -XX:+UseG1GC -XX:+UseContainerSupport -Xms256m -Xmx512m -XX:MaxDirectMemorySize=1536m -XX:+DisableExplicitGC -XX:InitiatingHeapOccupancyPercent=68 -Xlog:gc*=info:file=gc.log:time,tags:filecount=5,filesize=10M -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/app/log -XX:ErrorFile=/app/log/hs_err_pid%p.log"
 ENV LOG_LEVEL=INFO
 EXPOSE 8023
 ENTRYPOINT ${JAVA_HOME}/bin/java ${JAVA_OPTS} -Dsun.net.http.allowRestrictedHeaders=true --add-opens=java.base/java.nio=ALL-UNNAMED -jar ./dataproxy.jar

dataproxy-core/pom.xml

@@ -33,10 +33,10 @@
                     <groupId>org.apache.arrow</groupId>
                     <artifactId>arrow-format</artifactId>
                 </exclusion>
-                <exclusion>
+                <!--<exclusion>
                     <groupId>org.apache.arrow</groupId>
                     <artifactId>arrow-memory-core</artifactId>
-                </exclusion>
+                </exclusion>-->
                 <exclusion>
                     <groupId>org.apache.arrow</groupId>
                     <artifactId>arrow-memory</artifactId>
@@ -122,6 +122,12 @@
             </exclusions>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-memory-netty</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.github.ben-manes.caffeine</groupId>
             <artifactId>caffeine</artifactId>
@@ -146,5 +152,28 @@
             <artifactId>jsr305</artifactId>
         </dependency>
 
+        <!-- JUnit 5 -->
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <!-- Mockito for JUnit 5 -->
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-junit-jupiter</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/DefaultFlightServerConfigLoader.java

@@ -61,5 +61,6 @@ public void loadProperties(Properties properties) {
             throw new RuntimeException(e);
         }
         properties.put(FlightServerConfigKey.PORT, 8023);
+        properties.put(FlightServerConfigKey.METRICS_PORT, 9101);
     }
 }

dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfig.java

@@ -18,6 +18,9 @@
 
 import lombok.extern.slf4j.Slf4j;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.secretflow.dataproxy.core.listener.DataProxyAllocationListener;
 
 /**
  * @author yuexie
@@ -26,7 +29,14 @@
 @Slf4j
 public record FlightServerConfig(String host, int port) {
 
+    private static final BufferAllocator ROOT_ALLOCATOR =
+            new RootAllocator(new DataProxyAllocationListener(), 2L * 1024 * 1024 * 1024);
+
     public Location getLocation() {
         return Location.forGrpcInsecure(host, port);
     }
+
+    public BufferAllocator getBufferAllocator() {
+        return ROOT_ALLOCATOR;
+    }
 }

dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfigKey.java

@@ -26,4 +26,6 @@ public class FlightServerConfigKey {
 
     public static final String PORT = "SERVICE_PORT";
 
+    public static final String METRICS_PORT = "METRICS_PORT";
+
 }

dataproxy-core/src/main/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListener.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2025 Ant Group Co., Ltd.
+ *
+ * Licensed 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.secretflow.dataproxy.core.listener;
+
+import lombok.extern.slf4j.Slf4j;
+import org.apache.arrow.memory.AllocationListener;
+import org.apache.arrow.memory.AllocationOutcome;
+import org.apache.arrow.memory.BufferAllocator;
+
+/**
+ * @author yuexie
+ * @date 2025/4/14 16:14
+ **/
+@Slf4j
+public class DataProxyAllocationListener implements AllocationListener {
+
+    /**
+     * Called each time a new buffer has been requested.
+     *
+     * <p>An exception can be safely thrown by this method to terminate the allocation.
+     *
+     * @param size the buffer size being allocated
+     */
+    @Override
+    public void onPreAllocation(long size) {
+        AllocationListener.super.onPreAllocation(size);
+        log.debug("onPreAllocation, size: {}", size);
+    }
+
+    /**
+     * Called each time a new buffer has been allocated.
+     *
+     * <p>An exception cannot be thrown by this method.
+     *
+     * @param size the buffer size being allocated
+     */
+    @Override
+    public void onAllocation(long size) {
+        AllocationListener.super.onAllocation(size);
+        log.debug("onAllocation, size: {}", size);
+    }
+
+    /**
+     * Informed each time a buffer is released from allocation.
+     *
+     * <p>An exception cannot be thrown by this method.
+     *
+     * @param size The size of the buffer being released.
+     */
+    @Override
+    public void onRelease(long size) {
+        AllocationListener.super.onRelease(size);
+    }
+
+    /**
+     * Called whenever an allocation failed, giving the caller a chance to create some space in the
+     * allocator (either by freeing some resource, or by changing the limit), and, if successful,
+     * allowing the allocator to retry the allocation.
+     *
+     * @param size    the buffer size that was being allocated
+     * @param outcome the outcome of the failed allocation. Carries information of what failed
+     * @return true, if the allocation can be retried; false if the allocation should fail
+     */
+    @Override
+    public boolean onFailedAllocation(long size, AllocationOutcome outcome) {
+        log.debug("onFailedAllocation, size: {}, outcome: {}", size, outcome);
+        return AllocationListener.super.onFailedAllocation(size, outcome);
+    }
+
+    /**
+     * Called immediately after a child allocator was added to the parent allocator.
+     *
+     * @param parentAllocator The parent allocator to which a child was added
+     * @param childAllocator  The child allocator that was just added
+     */
+    @Override
+    public void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) {
+        AllocationListener.super.onChildAdded(parentAllocator, childAllocator);
+        log.debug("onChildAdded, childAllocator: {}, size: {}", childAllocator.getName(), childAllocator.getLimit());
+    }
+
+    /**
+     * Called immediately after a child allocator was removed from the parent allocator.
+     *
+     * @param parentAllocator The parent allocator from which a child was removed
+     * @param childAllocator  The child allocator that was just removed
+     */
+    @Override
+    public void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) {
+        AllocationListener.super.onChildRemoved(parentAllocator, childAllocator);
+        log.debug("onChildRemoved, childAllocator: {}, size: {}", childAllocator.getName(), childAllocator.getLimit());
+    }
+}

dataproxy-core/src/main/java/org/secretflow/dataproxy/core/reader/AbstractSender.java

@@ -81,8 +81,9 @@ public void send() {
                 ValueVectorUtility.ensureCapacity(root, takeRecordCount + 1);
                 this.toArrowVector(record, root, takeRecordCount);
                 takeRecordCount++;
-
-                if (takeRecordCount % 300_000 == 0) {
+                // 10w records, flush to arrow
+                // It can't be too big, and the off-heap memory is clipped
+                if (takeRecordCount % 100_000 == 0) {
                     break;
                 }
             }
@@ -158,7 +159,7 @@ public void close() throws Exception {
      * Pre-application for arrow vector memory
      */
     private void preAllocate() {
-
+        root.clear();
         ValueVectorUtility.preAllocate(root, estimatedRecordCount);
 
         root.getFieldVectors().forEach(fieldVector -> {
@@ -168,6 +169,5 @@ private void preAllocate() {
                 baseVariableWidthVector.allocateNew(estimatedRecordCount * 32);
             }
         });
-        root.clear();
     }
 }

dataproxy-core/src/test/java/org/secretflow/dataproxy/core/config/FlightServerContextTest.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2025 Ant Group Co., Ltd.
+ *
+ * Licensed 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.secretflow.dataproxy.core.config;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+/**
+ * @author yuexie
+ * @date 2025/06/16 10:59:10
+ */
+public class FlightServerContextTest {
+
+    @Test
+    void testGetInstanceInMultiThread() throws InterruptedException {
+        final int threadCount = 10;
+        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+        CountDownLatch latch = new CountDownLatch(threadCount);
+        AtomicReference<FlightServerContext> firstInstance = new AtomicReference<>();
+
+        for (int i = 0; i < threadCount; i++) {
+            executor.submit(() -> {
+                FlightServerContext instance = FlightServerContext.getInstance();
+                if (firstInstance.get() == null) {
+                    firstInstance.set(instance);
+                } else {
+                    assertSame(firstInstance.get(), instance, "The same instance should be returned in a multi-threaded environment");
+                }
+                latch.countDown();
+            });
+        }
+
+        latch.await(5, TimeUnit.SECONDS);
+        executor.shutdown();
+    }
+    @Test
+    void testGetOrDefault() {
+        String testKey = "nonexistent.key";
+        String defaultValue = "default";
+
+        String result = FlightServerContext.getOrDefault(testKey, String.class, defaultValue);
+        assertEquals(defaultValue, result);
+    }
+
+    @Test
+    void testFlightServerConfig() {
+        FlightServerContext context = FlightServerContext.getInstance();
+        assertNotNull(context.getFlightServerConfig(), "flightServerConfig should be initialized");
+    }
+
+}