Limit Edge query threads and remove unused binary allocator

integration-test/src/test/java/org/apache/iotdb/edge/it/IoTDBEdgeBasicIT.java

@@ -45,6 +45,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
+import java.nio.file.StandardOpenOption;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
@@ -53,6 +54,7 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashSet;
+import java.util.HexFormat;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -68,6 +70,7 @@
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -112,6 +115,11 @@ public static void setUp() throws Exception {
     ports = EnvUtils.searchAvailablePorts();
     rpcPort = ports[2];
     configurePorts(edgeHome.resolve("conf/iotdb-system.properties"));
+    // Existing installations may still carry this retired configuration key.
+    Files.writeString(
+        edgeHome.resolve("conf/iotdb-system.properties"),
+        "\nenable_binary_allocator=true\n",
+        StandardOpenOption.APPEND);
 
     runScript(edgeHome.resolve("sbin/start-edge.sh"), START_SCRIPT_LOG);
     edgePid = Long.parseLong(Files.readString(edgeHome.resolve("edge.pid")).trim());
@@ -246,6 +254,120 @@ private static Connection openTreeConnection() throws SQLException {
         jdbcUrl(), SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD);
   }
 
+  @Test
+  public void testConcurrentCrossDatabaseJoinsWithSmallDispatchPool() throws Exception {
+    try (Connection connection = openTableConnection();
+        Statement statement = connection.createStatement()) {
+      for (String database : new String[] {"edge_it_dispatch_a", "edge_it_dispatch_b"}) {
+        statement.execute("CREATE DATABASE " + database);
+        statement.execute("USE " + database);
+        statement.execute("CREATE TABLE sensor(device STRING TAG, value INT32 FIELD)");
+        statement.execute("INSERT INTO sensor(time,device,value) VALUES (1,'d1',42),(2,'d1',84)");
+      }
+    }
+    ExecutorService executor = Executors.newFixedThreadPool(4);
+    CountDownLatch ready = new CountDownLatch(4);
+    CountDownLatch start = new CountDownLatch(1);
+    List<Future<Void>> queries = new ArrayList<>();
+    try {
+      for (int i = 0; i < 4; i++) {
+        queries.add(
+            executor.submit(
+                () -> {
+                  try (Connection connection = openTableConnection();
+                      Statement statement = connection.createStatement()) {
+                    ready.countDown();
+                    assertTrue(start.await(30, TimeUnit.SECONDS));
+                    for (int iteration = 0; iteration < 10; iteration++) {
+                      try (ResultSet result =
+                          statement.executeQuery(
+                              "SELECT count(*), sum(a.value + b.value)"
+                                  + " FROM edge_it_dispatch_a.sensor a JOIN edge_it_dispatch_b.sensor b"
+                                  + " ON a.device = b.device AND a.time = b.time")) {
+                        assertTrue(result.next());
+                        assertEquals(2, result.getLong(1));
+                        assertEquals(252.0, result.getDouble(2), 0.0);
+                        assertFalse(result.next());
+                      }
+                    }
+                  }
+                  return null;
+                }));
+      }
+      assertTrue(ready.await(30, TimeUnit.SECONDS));
+      start.countDown();
+      for (Future<Void> query : queries) {
+        query.get(90, TimeUnit.SECONDS);
+      }
+    } finally {
+      start.countDown();
+      executor.shutdownNow();
+      assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    }
+    String threads = captureThreadDump("cross-database-joins");
+    long dispatchWorkers =
+        threads
+            .lines()
+            .filter(
+                line -> line.startsWith("\"pool-") && line.contains("Fragment-Instance-Dispatch-"))
+            .count();
+    assertTrue("The join must exercise fragment dispatch", dispatchWorkers > 0);
+    assertTrue("Fragment dispatch exceeded its configured pool size", dispatchWorkers <= 2);
+  }
+
+  @Test
+  public void testBlobReadWriteAfterLegacyAllocatorConfigurationReload() throws Exception {
+    byte[] expected = new byte[65536];
+    for (int i = 0; i < expected.length; i++) {
+      expected[i] = (byte) i;
+    }
+    try (Connection connection = openTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("LOAD CONFIGURATION");
+      statement.execute("CREATE DATABASE edge_it_blob");
+      statement.execute("USE edge_it_blob");
+      statement.execute("CREATE TABLE payloads(device STRING TAG, payload BLOB FIELD)");
+      statement.execute(
+          "INSERT INTO payloads(time,device,payload) VALUES (1,'d1',X'"
+              + HexFormat.of().formatHex(expected)
+              + "')");
+      for (int round = 0; round < 2; round++) {
+        try (ResultSet result = statement.executeQuery("SELECT payload FROM payloads")) {
+          assertTrue(result.next());
+          assertArrayEquals(expected, result.getBytes(1));
+          assertFalse(result.next());
+        }
+        if (round == 0) {
+          statement.execute("FLUSH");
+          statement.execute("LOAD CONFIGURATION");
+        }
+      }
+    }
+    assertFalse(captureThreadDump("blob-after-reload").contains("BinaryAllocator-"));
+  }
+
+  private static String captureThreadDump(String name) throws Exception {
+    Path output = WORK_DIR.resolve(name + "-threads.txt");
+    Process process =
+        new ProcessBuilder(
+                Paths.get(System.getProperty("java.home"), "bin", "jcmd").toString(),
+                Long.toString(edgePid),
+                "Thread.print",
+                "-l")
+            .redirectErrorStream(true)
+            .redirectOutput(output.toFile())
+            .start();
+    try {
+      assertTrue("Thread dump timed out", process.waitFor(30, TimeUnit.SECONDS));
+      assertEquals("Failed to capture thread dump: " + output, 0, process.exitValue());
+      return Files.readString(output);
+    } finally {
+      if (process.isAlive()) {
+        process.destroyForcibly();
+      }
+    }
+  }
+
   @Test
   public void testRatisMetadataConsensus() throws SQLException {
     Map<String, String> variables = new LinkedHashMap<>();
@@ -280,6 +402,8 @@ public void testPackagedConfiguration() throws Exception {
     assertEdgeProperty("coordinator_read_executor_size", "2");
     assertEdgeProperty("coordinator_scheduled_executor_size", "2");
     assertEdgeProperty("fragment_instance_notification_thread_count", "2");
+    assertEdgeProperty("driver_task_scheduler_notification_thread_count", "2");
+    assertEdgeProperty("fragment_instance_dispatch_thread_count", "2");
     assertEdgeProperty("cn_load_statistics_publisher_thread_count", "1");
     assertEdgeProperty("candidate_compaction_task_queue_size", "10");
     assertEdgeProperty("compaction_max_aligned_series_num_in_one_batch", "2");

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java

@@ -1026,6 +1026,12 @@ public class IoTDBConfig {
   /** Thread pool size for fragment instance state change notifications. */
   private int fragmentInstanceNotificationThreadCount = 4;
 
+  /** Zero retains the cached pool used by general deployments. */
+  private int driverTaskSchedulerNotificationThreadCount = 0;
+
+  /** Zero selects max(20, twice the available processors). */
+  private int fragmentInstanceDispatchThreadCount = 0;
+
   /** Policy of DataNodeSchemaCache eviction */
   private String dataNodeSchemaCacheEvictionPolicy = "FIFO";
 
@@ -3612,6 +3618,30 @@ public void setFragmentInstanceNotificationThreadCount(
     this.fragmentInstanceNotificationThreadCount = fragmentInstanceNotificationThreadCount;
   }
 
+  public int getDriverTaskSchedulerNotificationThreadCount() {
+    return driverTaskSchedulerNotificationThreadCount;
+  }
+
+  public void setDriverTaskSchedulerNotificationThreadCount(int threadCount) {
+    if (threadCount < 0) {
+      throw new IllegalArgumentException(
+          CommonMessages.EXCEPTION_THREAD_COUNT_MUST_BE_GREATER_THAN_OR_EQUAL_TO_0_988EF69B);
+    }
+    this.driverTaskSchedulerNotificationThreadCount = threadCount;
+  }
+
+  public int getFragmentInstanceDispatchThreadCount() {
+    return fragmentInstanceDispatchThreadCount;
+  }
+
+  public void setFragmentInstanceDispatchThreadCount(int threadCount) {
+    if (threadCount < 0) {
+      throw new IllegalArgumentException(
+          CommonMessages.EXCEPTION_THREAD_COUNT_MUST_BE_GREATER_THAN_OR_EQUAL_TO_0_988EF69B);
+    }
+    this.fragmentInstanceDispatchThreadCount = threadCount;
+  }
+
   public TEndPoint getAddressAndPort() {
     return new TEndPoint(rpcAddress, rpcPort);
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -19,7 +19,6 @@
 package org.apache.iotdb.db.conf;
 
 import org.apache.iotdb.calc.exception.QueryProcessException;
-import org.apache.iotdb.commons.binaryallocator.BinaryAllocator;
 import org.apache.iotdb.commons.conf.CommonConfig;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.conf.ConfigurationFileUtils;
@@ -1012,6 +1011,16 @@ public void loadProperties(TrimProperties properties) throws BadNodeUrlException
             properties.getProperty(
                 "fragment_instance_notification_thread_count",
                 Integer.toString(conf.getFragmentInstanceNotificationThreadCount()))));
+    conf.setDriverTaskSchedulerNotificationThreadCount(
+        Integer.parseInt(
+            properties.getProperty(
+                "driver_task_scheduler_notification_thread_count",
+                Integer.toString(conf.getDriverTaskSchedulerNotificationThreadCount()))));
+    conf.setFragmentInstanceDispatchThreadCount(
+        Integer.parseInt(
+            properties.getProperty(
+                "fragment_instance_dispatch_thread_count",
+                Integer.toString(conf.getFragmentInstanceDispatchThreadCount()))));
     conf.setDataNodeTableSchemaCacheSize(
         Long.parseLong(
             properties.getProperty(
@@ -2308,21 +2317,6 @@ public synchronized void loadHotModifiedProps(TrimProperties properties)
       // update retry config
       commonDescriptor.loadRetryProperties(properties);
 
-      // update binary allocator
-      commonDescriptor
-          .getConfig()
-          .setEnableBinaryAllocator(
-              Boolean.parseBoolean(
-                  properties.getProperty(
-                      "enable_binary_allocator",
-                      ConfigurationFileUtils.getConfigurationDefaultValue(
-                          "enable_binary_allocator"))));
-      if (commonDescriptor.getConfig().isEnableBinaryAllocator()) {
-        BinaryAllocator.getInstance().start();
-      } else {
-        BinaryAllocator.getInstance().close(true);
-      }
-
       // update disk_space_warning_threshold; also refresh the static copy in JVMCommonUtils that
       // the ReadOnly disk guard reads, otherwise the new threshold would not take effect until
       // restart. Parsing / validation is shared with the ConfigNode hot-reload path.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskThread.java

@@ -23,6 +23,7 @@
 import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
 import org.apache.iotdb.commons.concurrent.ThreadName;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.queryengine.execution.driver.IDriver;
 import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.MultilevelPriorityQueue;
 import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTask;
@@ -33,6 +34,7 @@
 import io.airlift.units.Duration;
 
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.IntStream;
 
@@ -54,10 +56,13 @@ public class DriverTaskThread extends AbstractDriverThread {
                       (level + 1) * DRIVER_TASK_EXECUTION_TIME_SLICE_IN_MS, TimeUnit.MILLISECONDS))
           .toArray(Duration[]::new);
 
-  // We manage thread pool size directly, so create an unlimited pool
-  private static final Executor listeningExecutor =
-      IoTDBThreadPoolFactory.newCachedThreadPool(
-          ThreadName.DRIVER_TASK_SCHEDULER_NOTIFICATION.getName());
+  private static final Executor NOTIFICATION_EXECUTOR =
+      createNotificationExecutor(
+          IoTDBDescriptor.getInstance()
+              .getConfig()
+              .getDriverTaskSchedulerNotificationThreadCount());
+
+  private final Executor listeningExecutor;
 
   private final Ticker ticker;
 
@@ -67,10 +72,32 @@ public DriverTaskThread(
       IndexedBlockingQueue<DriverTask> queue,
       ITaskScheduler scheduler,
       ThreadProducer producer) {
+    this(workerId, tg, queue, scheduler, producer, NOTIFICATION_EXECUTOR);
+  }
+
+  DriverTaskThread(
+      String workerId,
+      ThreadGroup tg,
+      IndexedBlockingQueue<DriverTask> queue,
+      ITaskScheduler scheduler,
+      ThreadProducer producer,
+      Executor listeningExecutor) {
     super(workerId, tg, queue, scheduler, producer);
+    this.listeningExecutor = listeningExecutor;
     this.ticker = Ticker.systemTicker();
   }
 
+  static ExecutorService createNotificationExecutor(int threadCount) {
+    String poolName = ThreadName.DRIVER_TASK_SCHEDULER_NOTIFICATION.getName();
+    if (threadCount == 0) {
+      return IoTDBThreadPoolFactory.newCachedThreadPool(poolName);
+    }
+    // Queue notifications instead of running them on threads completing driver futures, which
+    // may hold locks. Idle workers can exit after the query workload ends.
+    return IoTDBThreadPoolFactory.newFixedThreadPoolWithIdleThreadTimeout(
+        threadCount, 60, TimeUnit.SECONDS, poolName);
+  }
+
   @Override
   public void execute(DriverTask task) throws InterruptedException {
     long startNanos = ticker.read();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java

@@ -284,7 +284,10 @@ private Coordinator() {
     this.typeManager = new InternalTypeManager();
     this.executor = getQueryExecutor();
     this.scheduledExecutor = getScheduledExecutor();
-    int dispatchThreadNum = Math.max(20, Runtime.getRuntime().availableProcessors() * 2);
+    int dispatchThreadNum = CONFIG.getFragmentInstanceDispatchThreadCount();
+    if (dispatchThreadNum == 0) {
+      dispatchThreadNum = Math.max(20, Runtime.getRuntime().availableProcessors() * 2);
+    }
     this.dispatchExecutor =
         IoTDBThreadPoolFactory.newCachedThreadPool(
             ThreadName.FRAGMENT_INSTANCE_DISPATCH.getName(),

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/QueryThreadPoolConfigTest.java

@@ -32,6 +32,15 @@ public class QueryThreadPoolConfigTest {
   public void testDefaultsAndPositiveSizes() throws Exception {
     IoTDBConfig config = new IoTDBConfig();
     assertSizes(config, 20, 10, 4);
+    assertAdditionalSizes(config, 0, 0);
+    assertEquals(
+        "0",
+        ConfigurationFileUtils.getConfigurationDefaultValue(
+            "driver_task_scheduler_notification_thread_count"));
+    assertEquals(
+        "0",
+        ConfigurationFileUtils.getConfigurationDefaultValue(
+            "fragment_instance_dispatch_thread_count"));
     assertEquals(
         "20",
         ConfigurationFileUtils.getConfigurationDefaultValue("coordinator_read_executor_size"));
@@ -54,6 +63,15 @@ public void testDefaultsAndPositiveSizes() throws Exception {
           () -> config.setFragmentInstanceNotificationThreadCount(invalid));
     }
     assertSizes(config, 20, 10, 4);
+    for (int invalid : new int[] {-1, Integer.MIN_VALUE}) {
+      assertThrows(
+          IllegalArgumentException.class,
+          () -> config.setDriverTaskSchedulerNotificationThreadCount(invalid));
+      assertThrows(
+          IllegalArgumentException.class,
+          () -> config.setFragmentInstanceDispatchThreadCount(invalid));
+    }
+    assertAdditionalSizes(config, 0, 0);
   }
 
   @Test
@@ -63,14 +81,25 @@ public void testStartupOverridesAreRestartOnly() throws Exception {
     properties.setProperty("coordinator_read_executor_size", "3");
     properties.setProperty("coordinator_scheduled_executor_size", "2");
     properties.setProperty("fragment_instance_notification_thread_count", "1");
+    properties.setProperty("driver_task_scheduler_notification_thread_count", "1");
+    properties.setProperty("fragment_instance_dispatch_thread_count", "2");
     descriptor.loadProperties(properties);
     assertSizes(descriptor.getConfig(), 3, 2, 1);
+    assertAdditionalSizes(descriptor.getConfig(), 1, 2);
 
     properties.setProperty("coordinator_read_executor_size", "6");
     properties.setProperty("coordinator_scheduled_executor_size", "5");
     properties.setProperty("fragment_instance_notification_thread_count", "4");
+    properties.setProperty("driver_task_scheduler_notification_thread_count", "4");
+    properties.setProperty("fragment_instance_dispatch_thread_count", "5");
     descriptor.loadHotModifiedProps(properties);
     assertSizes(descriptor.getConfig(), 3, 2, 1);
+    assertAdditionalSizes(descriptor.getConfig(), 1, 2);
+  }
+
+  private static void assertAdditionalSizes(IoTDBConfig config, int notification, int dispatch) {
+    assertEquals(notification, config.getDriverTaskSchedulerNotificationThreadCount());
+    assertEquals(dispatch, config.getFragmentInstanceDispatchThreadCount());
   }
 
   private static void assertSizes(IoTDBConfig config, int read, int scheduled, int notification) {