perf: 优化多源联邦分片调度
- 增加 Engine 级有界并行调度与资源收口 - 补齐调度饱和、失败快速传播及六源基准测试
This commit is contained in:
@@ -0,0 +1,129 @@
|
|||||||
|
package com.easyagents.federation.sql.adapter.jdbc;
|
||||||
|
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlEngine;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlEngines;
|
||||||
|
import com.easyagents.federation.sql.api.SqlQueryCommand;
|
||||||
|
import com.easyagents.federation.sql.execute.FederationResultCursor;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
|
||||||
|
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
|
||||||
|
import com.easyagents.federation.sql.source.FederationSourceDefinition;
|
||||||
|
import com.easyagents.federation.sql.source.RuntimeFingerprint;
|
||||||
|
import com.easyagents.federation.sql.source.SourceApplyOptions;
|
||||||
|
import com.easyagents.federation.sql.source.SourceId;
|
||||||
|
import com.mysql.cj.jdbc.MysqlDataSource;
|
||||||
|
import org.junit.Assert;
|
||||||
|
import org.junit.Assume;
|
||||||
|
import org.junit.Test;
|
||||||
|
|
||||||
|
import java.math.BigDecimal;
|
||||||
|
import java.sql.Connection;
|
||||||
|
import java.sql.PreparedStatement;
|
||||||
|
import java.sql.ResultSet;
|
||||||
|
import java.sql.SQLException;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.Map;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 亿级数据集 100,000 行 Cursor 手工基准;默认跳过,必须显式启用。
|
||||||
|
*/
|
||||||
|
public class Jdbc100mCursorManualBenchmarkTest {
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 从支付库顺序消费恰好 100,000 行,验证 Cursor 流式消费和资源上限。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldConsumeOneHundredThousandRows() {
|
||||||
|
Assume.assumeTrue(Boolean.getBoolean("federation.100m.enabled"));
|
||||||
|
String host = System.getProperty("federation.100m.mysql.host", "127.0.0.1");
|
||||||
|
int port = Integer.getInteger("federation.100m.mysql.port", 33307);
|
||||||
|
String database = System.getProperty(
|
||||||
|
"federation.100m.mysql.database", "ef_bank_payment_perf");
|
||||||
|
String username = System.getProperty("federation.100m.mysql.username", "ef_bench_ro");
|
||||||
|
String password = System.getProperty("federation.100m.mysql.password");
|
||||||
|
Assert.assertNotNull("必须通过系统属性提供只读密码", password);
|
||||||
|
|
||||||
|
MysqlDataSource dataSource = new MysqlDataSource();
|
||||||
|
dataSource.setUrl("jdbc:mysql://" + host + ':' + port + '/' + database
|
||||||
|
+ "?useSSL=false&allowPublicKeyRetrieval=true&serverTimezone=Asia/Shanghai"
|
||||||
|
+ "&useCursorFetch=true&defaultFetchSize=1000");
|
||||||
|
dataSource.setUser(username);
|
||||||
|
dataSource.setPassword(password);
|
||||||
|
|
||||||
|
SourceId sourceId = new SourceId("payment-100m");
|
||||||
|
FederationSourceDefinition definition = new FederationSourceDefinition(
|
||||||
|
sourceId, 1, JdbcFederationSqlAdapterProvider.ADAPTER_ID,
|
||||||
|
List.of(new JdbcSchemaDefinition("MAIN", database, null)), Map.of());
|
||||||
|
FederationExecutionPolicy policy = new FederationExecutionPolicy(
|
||||||
|
6, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);
|
||||||
|
FederationQueryScopeDefinition scope = FederationQueryScopeDefinition.virtual(
|
||||||
|
"payment-cursor-100m", 1,
|
||||||
|
Map.of("PAY", FederationSourceBindingDefinition.of(sourceId, 1)),
|
||||||
|
"PAY", policy);
|
||||||
|
|
||||||
|
CursorResult nativeResult = consumeNative(dataSource);
|
||||||
|
CursorResult federationResult;
|
||||||
|
long federationStarted = System.nanoTime();
|
||||||
|
try (FederationSqlEngine engine = FederationSqlEngines.builder()
|
||||||
|
.dataSourceResolver(ignored -> FederationDataSourceHandles.shared(
|
||||||
|
dataSource, new RuntimeFingerprint("MySQL", "8.0", "JDBC", "8.4", "1")))
|
||||||
|
.adapter(new JdbcFederationSqlAdapterProvider())
|
||||||
|
.federationExecutionPolicy(policy)
|
||||||
|
.build()) {
|
||||||
|
engine.sources().apply(definition, SourceApplyOptions.prewarmNow());
|
||||||
|
String sql = "SELECT payment_id, amount FROM PAY.MAIN.payment_order "
|
||||||
|
+ "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
|
||||||
|
int rows = 0;
|
||||||
|
long amountCents = 0;
|
||||||
|
try (FederationResultCursor cursor = engine.query(
|
||||||
|
SqlQueryCommand.of(sql, scope, List.of()))) {
|
||||||
|
while (cursor.next()) {
|
||||||
|
rows++;
|
||||||
|
amountCents += ((BigDecimal) cursor.getObject(2))
|
||||||
|
.movePointRight(2).longValue();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
federationResult = new CursorResult(rows, amountCents,
|
||||||
|
elapsedMillis(federationStarted));
|
||||||
|
}
|
||||||
|
Assert.assertEquals(100_000, nativeResult.rows());
|
||||||
|
Assert.assertEquals(100_000, federationResult.rows());
|
||||||
|
Assert.assertEquals(nativeResult.amountCents(), federationResult.amountCents());
|
||||||
|
Assert.assertTrue(federationResult.amountCents() > 0L);
|
||||||
|
System.out.println("BENCH_CURSOR_100K rows=" + federationResult.rows()
|
||||||
|
+ " nativeMillis=" + nativeResult.elapsedMillis()
|
||||||
|
+ " federationMillis=" + federationResult.elapsedMillis()
|
||||||
|
+ " amountCents=" + federationResult.amountCents());
|
||||||
|
}
|
||||||
|
|
||||||
|
private CursorResult consumeNative(MysqlDataSource dataSource) {
|
||||||
|
String sql = "SELECT payment_id, amount FROM payment_order "
|
||||||
|
+ "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
|
||||||
|
long started = System.nanoTime();
|
||||||
|
int rows = 0;
|
||||||
|
long amountCents = 0;
|
||||||
|
try (Connection connection = dataSource.getConnection();
|
||||||
|
PreparedStatement statement = connection.prepareStatement(sql,
|
||||||
|
ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
|
||||||
|
statement.setFetchSize(1000);
|
||||||
|
try (ResultSet resultSet = statement.executeQuery()) {
|
||||||
|
while (resultSet.next()) {
|
||||||
|
rows++;
|
||||||
|
amountCents += resultSet.getBigDecimal(2)
|
||||||
|
.movePointRight(2).longValue();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} catch (SQLException exception) {
|
||||||
|
throw new IllegalStateException("原生 JDBC Cursor 基准执行失败", exception);
|
||||||
|
}
|
||||||
|
return new CursorResult(rows, amountCents, elapsedMillis(started));
|
||||||
|
}
|
||||||
|
|
||||||
|
private long elapsedMillis(long started) {
|
||||||
|
return (System.nanoTime() - started) / 1_000_000L;
|
||||||
|
}
|
||||||
|
|
||||||
|
private record CursorResult(int rows, long amountCents, long elapsedMillis) {
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -27,6 +27,7 @@ import com.easyagents.federation.sql.source.RuntimeFingerprint;
|
|||||||
import com.easyagents.federation.sql.source.SourceApplyOptions;
|
import com.easyagents.federation.sql.source.SourceApplyOptions;
|
||||||
import com.easyagents.federation.sql.source.SourceId;
|
import com.easyagents.federation.sql.source.SourceId;
|
||||||
import java.math.BigDecimal;
|
import java.math.BigDecimal;
|
||||||
|
import java.sql.SQLException;
|
||||||
import java.time.Instant;
|
import java.time.Instant;
|
||||||
import java.sql.Connection;
|
import java.sql.Connection;
|
||||||
import java.sql.Statement;
|
import java.sql.Statement;
|
||||||
@@ -34,6 +35,9 @@ import java.sql.Types;
|
|||||||
import java.util.ArrayList;
|
import java.util.ArrayList;
|
||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.Map;
|
import java.util.Map;
|
||||||
|
import java.util.concurrent.CountDownLatch;
|
||||||
|
import java.util.concurrent.TimeUnit;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
import java.util.concurrent.atomic.AtomicLong;
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
import java.util.concurrent.atomic.AtomicReference;
|
import java.util.concurrent.atomic.AtomicReference;
|
||||||
import org.h2.jdbcx.JdbcDataSource;
|
import org.h2.jdbcx.JdbcDataSource;
|
||||||
@@ -50,6 +54,13 @@ public class JdbcFederatedQueryEngineTest {
|
|||||||
private static final SourceId SALES_SOURCE = new SourceId("sales-source");
|
private static final SourceId SALES_SOURCE = new SourceId("sales-source");
|
||||||
private static final SourceId BILLING_SOURCE = new SourceId("billing-source");
|
private static final SourceId BILLING_SOURCE = new SourceId("billing-source");
|
||||||
private static final SourceId REGION_SOURCE = new SourceId("region-source");
|
private static final SourceId REGION_SOURCE = new SourceId("region-source");
|
||||||
|
private static final AtomicInteger ACTIVE_FRAGMENTS = new AtomicInteger();
|
||||||
|
private static final AtomicInteger MAXIMUM_ACTIVE_FRAGMENTS = new AtomicInteger();
|
||||||
|
private static final AtomicInteger FAILING_TRACKED_VALUE = new AtomicInteger(
|
||||||
|
Integer.MIN_VALUE
|
||||||
|
);
|
||||||
|
private static final AtomicReference<CountDownLatch> FRAGMENT_START_BARRIER =
|
||||||
|
new AtomicReference<>(new CountDownLatch(0));
|
||||||
|
|
||||||
private JdbcDataSource sales;
|
private JdbcDataSource sales;
|
||||||
private JdbcDataSource billing;
|
private JdbcDataSource billing;
|
||||||
@@ -80,7 +91,10 @@ public class JdbcFederatedQueryEngineTest {
|
|||||||
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 12:00:00.123456+08:00')",
|
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 12:00:00.123456+08:00')",
|
||||||
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
|
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
|
||||||
"INSERT INTO PRECISE_EVENT VALUES "
|
"INSERT INTO PRECISE_EVENT VALUES "
|
||||||
+ "(1, TIMESTAMP '2026-08-21 12:00:00.123456')");
|
+ "(1, TIMESTAMP '2026-08-21 12:00:00.123456')",
|
||||||
|
fragmentTrackingAlias(),
|
||||||
|
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(ID) AS TRACKED_KEY "
|
||||||
|
+ "FROM CUSTOMER WHERE ID <= 2");
|
||||||
execute(billing,
|
execute(billing,
|
||||||
"CREATE TABLE ORDER_ITEM (ID INT PRIMARY KEY, CUSTOMER_ID INT NOT NULL, AMOUNT DECIMAL(12,2), CODE VARCHAR(64))",
|
"CREATE TABLE ORDER_ITEM (ID INT PRIMARY KEY, CUSTOMER_ID INT NOT NULL, AMOUNT DECIMAL(12,2), CODE VARCHAR(64))",
|
||||||
"INSERT INTO ORDER_ITEM VALUES (10, 1, 30.00, 'Alice'), (11, 1, 20.00, 'Alice'), (12, 2, 80.00, 'Bob')",
|
"INSERT INTO ORDER_ITEM VALUES (10, 1, 30.00, 'Alice'), (11, 1, 20.00, 'Alice'), (12, 2, 80.00, 'Bob')",
|
||||||
@@ -90,10 +104,16 @@ public class JdbcFederatedQueryEngineTest {
|
|||||||
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 06:00:00.123456+02:00')",
|
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 06:00:00.123456+02:00')",
|
||||||
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
|
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
|
||||||
"INSERT INTO PRECISE_EVENT VALUES "
|
"INSERT INTO PRECISE_EVENT VALUES "
|
||||||
+ "(2, TIMESTAMP '2026-08-21 08:30:00.654321')");
|
+ "(2, TIMESTAMP '2026-08-21 08:30:00.654321')",
|
||||||
|
fragmentTrackingAlias(),
|
||||||
|
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(ID) AS TRACKED_KEY "
|
||||||
|
+ "FROM ORDER_ITEM WHERE ID <= 11");
|
||||||
execute(region,
|
execute(region,
|
||||||
"CREATE TABLE CUSTOMER_REGION (CUSTOMER_ID INT PRIMARY KEY, REGION VARCHAR(64))",
|
"CREATE TABLE CUSTOMER_REGION (CUSTOMER_ID INT PRIMARY KEY, REGION VARCHAR(64))",
|
||||||
"INSERT INTO CUSTOMER_REGION VALUES (1, 'North'), (2, 'South'), (3, 'West')");
|
"INSERT INTO CUSTOMER_REGION VALUES (1, 'North'), (2, 'South'), (3, 'West')",
|
||||||
|
fragmentTrackingAlias(),
|
||||||
|
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(CUSTOMER_ID) AS TRACKED_KEY "
|
||||||
|
+ "FROM CUSTOMER_REGION WHERE CUSTOMER_ID <= 2");
|
||||||
|
|
||||||
RuntimeFingerprint fingerprint = new RuntimeFingerprint(
|
RuntimeFingerprint fingerprint = new RuntimeFingerprint(
|
||||||
"H2", "2", "H2 JDBC Driver", "2", "1"
|
"H2", "2", "H2 JDBC Driver", "2", "1"
|
||||||
@@ -356,6 +376,91 @@ public class JdbcFederatedQueryEngineTest {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证真实 JDBC Fragment 执行会重叠,且始终受查询级并发上限约束。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldBoundConcurrentJdbcFragments() {
|
||||||
|
ACTIVE_FRAGMENTS.set(0);
|
||||||
|
MAXIMUM_ACTIVE_FRAGMENTS.set(0);
|
||||||
|
FRAGMENT_START_BARRIER.set(new CountDownLatch(2));
|
||||||
|
FederationQueryScopeDefinition threeSourceScope =
|
||||||
|
new FederationQueryScopeDefinition(
|
||||||
|
"tracked-fragments",
|
||||||
|
1,
|
||||||
|
Map.of(
|
||||||
|
"SALES", FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
|
||||||
|
"BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1),
|
||||||
|
"REGION", FederationSourceBindingDefinition.of(REGION_SOURCE, 1)
|
||||||
|
),
|
||||||
|
"SALES",
|
||||||
|
threeSourcePolicy()
|
||||||
|
);
|
||||||
|
|
||||||
|
try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
|
||||||
|
"SELECT COUNT(*) FROM SALES.APP.TRACKED_VALUE s "
|
||||||
|
+ "JOIN BILLING.APP.TRACKED_VALUE b ON s.TRACKED_KEY = b.TRACKED_KEY "
|
||||||
|
+ "JOIN REGION.APP.TRACKED_VALUE r ON s.TRACKED_KEY = r.TRACKED_KEY",
|
||||||
|
threeSourceScope,
|
||||||
|
List.of()
|
||||||
|
))) {
|
||||||
|
Assert.assertTrue(cursor.next());
|
||||||
|
Assert.assertEquals(8L, cursor.getObject(1));
|
||||||
|
Assert.assertFalse(cursor.next());
|
||||||
|
}
|
||||||
|
|
||||||
|
Assert.assertEquals(2, MAXIMUM_ACTIVE_FRAGMENTS.get());
|
||||||
|
Assert.assertEquals(0, ACTIVE_FRAGMENTS.get());
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证任一并行 Fragment 失败后立即传播原始执行错误,不等待未启动任务超时。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldFailFastWhenConcurrentFragmentFails() {
|
||||||
|
ACTIVE_FRAGMENTS.set(0);
|
||||||
|
MAXIMUM_ACTIVE_FRAGMENTS.set(0);
|
||||||
|
FRAGMENT_START_BARRIER.set(new CountDownLatch(0));
|
||||||
|
FAILING_TRACKED_VALUE.set(10);
|
||||||
|
FederationQueryScopeDefinition threeSourceScope =
|
||||||
|
new FederationQueryScopeDefinition(
|
||||||
|
"tracked-fragment-failure",
|
||||||
|
1,
|
||||||
|
Map.of(
|
||||||
|
"SALES", FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
|
||||||
|
"BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1),
|
||||||
|
"REGION", FederationSourceBindingDefinition.of(REGION_SOURCE, 1)
|
||||||
|
),
|
||||||
|
"SALES",
|
||||||
|
new FederationExecutionPolicy(
|
||||||
|
3,
|
||||||
|
8,
|
||||||
|
2,
|
||||||
|
100_000,
|
||||||
|
64L * 1024L * 1024L,
|
||||||
|
1_000
|
||||||
|
)
|
||||||
|
);
|
||||||
|
|
||||||
|
try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
|
||||||
|
"SELECT COUNT(*) FROM SALES.APP.TRACKED_VALUE s "
|
||||||
|
+ "JOIN BILLING.APP.TRACKED_VALUE b ON s.TRACKED_KEY = b.TRACKED_KEY "
|
||||||
|
+ "JOIN REGION.APP.TRACKED_VALUE r ON s.TRACKED_KEY = r.TRACKED_KEY",
|
||||||
|
threeSourceScope,
|
||||||
|
List.of()
|
||||||
|
))) {
|
||||||
|
cursor.next();
|
||||||
|
Assert.fail("failing fragment should terminate the federated query");
|
||||||
|
} catch (FederationSqlException exception) {
|
||||||
|
Assert.assertEquals(
|
||||||
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
|
exception.errorCode()
|
||||||
|
);
|
||||||
|
} finally {
|
||||||
|
FAILING_TRACKED_VALUE.set(Integer.MIN_VALUE);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 验证两个分片分别绑定原始查询中的动态参数。
|
* 验证两个分片分别绑定原始查询中的动态参数。
|
||||||
*/
|
*/
|
||||||
@@ -717,6 +822,41 @@ public class JdbcFederatedQueryEngineTest {
|
|||||||
return dataSource;
|
return dataSource;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* H2 视图调用的并发跟踪函数。
|
||||||
|
*
|
||||||
|
* @param value 原始列值
|
||||||
|
* @return 固定键值
|
||||||
|
* @throws SQLException 分片未能在时限内重叠启动
|
||||||
|
*/
|
||||||
|
public static int trackFragment(int value) throws SQLException {
|
||||||
|
int active = ACTIVE_FRAGMENTS.incrementAndGet();
|
||||||
|
MAXIMUM_ACTIVE_FRAGMENTS.accumulateAndGet(active, Math::max);
|
||||||
|
CountDownLatch barrier = FRAGMENT_START_BARRIER.get();
|
||||||
|
barrier.countDown();
|
||||||
|
try {
|
||||||
|
if (value == FAILING_TRACKED_VALUE.get()) {
|
||||||
|
throw new SQLException("simulated federation fragment failure");
|
||||||
|
}
|
||||||
|
if (!barrier.await(2, TimeUnit.SECONDS)) {
|
||||||
|
throw new SQLException("federation fragments did not overlap");
|
||||||
|
}
|
||||||
|
Thread.sleep(100L);
|
||||||
|
return 1;
|
||||||
|
} catch (InterruptedException exception) {
|
||||||
|
Thread.currentThread().interrupt();
|
||||||
|
throw new SQLException("fragment tracking was interrupted", exception);
|
||||||
|
} finally {
|
||||||
|
ACTIVE_FRAGMENTS.decrementAndGet();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String fragmentTrackingAlias() {
|
||||||
|
return "CREATE ALIAS TRACK_FRAGMENT FOR \""
|
||||||
|
+ JdbcFederatedQueryEngineTest.class.getName()
|
||||||
|
+ ".trackFragment\"";
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 根据物理源选择测试数据库。
|
* 根据物理源选择测试数据库。
|
||||||
*
|
*
|
||||||
|
|||||||
@@ -0,0 +1,185 @@
|
|||||||
|
package com.easyagents.federation.sql.adapter.jdbc;
|
||||||
|
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlEngine;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlEngines;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlException;
|
||||||
|
import com.easyagents.federation.sql.api.SqlQueryCommand;
|
||||||
|
import com.easyagents.federation.sql.compile.FederationSqlPlan;
|
||||||
|
import com.easyagents.federation.sql.compile.SqlCompileRequest;
|
||||||
|
import com.easyagents.federation.sql.execute.FederationResultCursor;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationQueryMode;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
|
||||||
|
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
|
||||||
|
import com.easyagents.federation.sql.source.FederationSourceDefinition;
|
||||||
|
import com.easyagents.federation.sql.source.RuntimeFingerprint;
|
||||||
|
import com.easyagents.federation.sql.source.SourceApplyOptions;
|
||||||
|
import com.easyagents.federation.sql.source.SourceId;
|
||||||
|
import org.h2.jdbcx.JdbcDataSource;
|
||||||
|
import org.junit.After;
|
||||||
|
import org.junit.Assert;
|
||||||
|
import org.junit.Before;
|
||||||
|
import org.junit.Test;
|
||||||
|
|
||||||
|
import java.sql.Connection;
|
||||||
|
import java.sql.Statement;
|
||||||
|
import java.util.LinkedHashMap;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.Map;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 四源和六源 JDBC 多阶段 Join 集成测试。
|
||||||
|
*/
|
||||||
|
public class JdbcSixSourceQueryEngineTest {
|
||||||
|
|
||||||
|
private static final FederationExecutionPolicy SIX_SOURCE_POLICY =
|
||||||
|
new FederationExecutionPolicy(6, 12, 2, 100_000,
|
||||||
|
64L * 1024L * 1024L, 60_000);
|
||||||
|
|
||||||
|
private final Map<SourceId, JdbcDataSource> dataSources = new LinkedHashMap<>();
|
||||||
|
private final Map<String, FederationSourceBindingDefinition> bindings = new LinkedHashMap<>();
|
||||||
|
private FederationSqlEngine engine;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 创建六个独立 H2 JDBC 数据源。
|
||||||
|
*/
|
||||||
|
@Before
|
||||||
|
public void setUp() throws Exception {
|
||||||
|
for (int index = 1; index <= 6; index++) {
|
||||||
|
SourceId sourceId = new SourceId("source-" + index);
|
||||||
|
JdbcDataSource dataSource = dataSource("six-source-" + index);
|
||||||
|
execute(dataSource,
|
||||||
|
"DROP TABLE IF EXISTS SUBJECT_METRIC",
|
||||||
|
"CREATE TABLE SUBJECT_METRIC (SUBJECT_ID INT PRIMARY KEY, METRIC_VALUE INT NOT NULL)",
|
||||||
|
"INSERT INTO SUBJECT_METRIC VALUES (1, " + (index * 10)
|
||||||
|
+ "), (2, " + (index * 100) + ")");
|
||||||
|
dataSources.put(sourceId, dataSource);
|
||||||
|
bindings.put("S" + index, FederationSourceBindingDefinition.of(sourceId, 1));
|
||||||
|
}
|
||||||
|
RuntimeFingerprint fingerprint = new RuntimeFingerprint(
|
||||||
|
"H2", "2", "H2 JDBC Driver", "2", "1");
|
||||||
|
engine = FederationSqlEngines.builder()
|
||||||
|
.dataSourceResolver(definition -> FederationDataSourceHandles.shared(
|
||||||
|
dataSources.get(definition.sourceId()), fingerprint))
|
||||||
|
.adapter(new JdbcFederationSqlAdapterProvider())
|
||||||
|
.federationExecutionPolicy(SIX_SOURCE_POLICY)
|
||||||
|
.maximumPlanCacheEntries(32)
|
||||||
|
.build();
|
||||||
|
for (SourceId sourceId : dataSources.keySet()) {
|
||||||
|
engine.sources().apply(definition(sourceId), SourceApplyOptions.prewarmNow());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 关闭 Engine。
|
||||||
|
*/
|
||||||
|
@After
|
||||||
|
public void tearDown() {
|
||||||
|
if (engine != null) {
|
||||||
|
engine.close();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证同一六源 Scope 中只引用四源时生成四个 Fragment。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldCompileFourReferencedSources() {
|
||||||
|
FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
|
||||||
|
selectSql(4), scope(SIX_SOURCE_POLICY)));
|
||||||
|
|
||||||
|
Assert.assertEquals(FederationQueryMode.FEDERATED, plan.queryMode());
|
||||||
|
Assert.assertEquals(4, plan.referencedSources().size());
|
||||||
|
Assert.assertEquals(4, plan.fragments().size());
|
||||||
|
Assert.assertEquals(3, plan.joinOptimizations().size());
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证六个独立 JDBC 数据源可完成五阶段 Join 并返回稳定结果。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldExecuteJoinAcrossSixSources() {
|
||||||
|
String sql = selectSql(6);
|
||||||
|
FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
|
||||||
|
sql, scope(SIX_SOURCE_POLICY)));
|
||||||
|
|
||||||
|
Assert.assertEquals(6, plan.referencedSources().size());
|
||||||
|
Assert.assertEquals(6, plan.fragments().size());
|
||||||
|
Assert.assertEquals(5, plan.joinOptimizations().size());
|
||||||
|
|
||||||
|
try (FederationResultCursor cursor = engine.query(
|
||||||
|
SqlQueryCommand.of(sql, scope(SIX_SOURCE_POLICY), List.of()))) {
|
||||||
|
Assert.assertTrue(cursor.next());
|
||||||
|
Assert.assertEquals(1, cursor.getObject(1));
|
||||||
|
Assert.assertEquals(10, cursor.getObject(2));
|
||||||
|
Assert.assertEquals(60, cursor.getObject(7));
|
||||||
|
Assert.assertTrue(cursor.next());
|
||||||
|
Assert.assertEquals(2, cursor.getObject(1));
|
||||||
|
Assert.assertEquals(600, cursor.getObject(7));
|
||||||
|
Assert.assertFalse(cursor.next());
|
||||||
|
Assert.assertEquals(6, cursor.metrics().fragments().size());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证 Scope 收紧为四源时,六源 SQL 在编译期稳定拒绝。
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldRejectSixSourcesWhenScopeLimitIsFour() {
|
||||||
|
FederationExecutionPolicy fourSourcePolicy = new FederationExecutionPolicy(
|
||||||
|
4, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);
|
||||||
|
|
||||||
|
try {
|
||||||
|
engine.compile(SqlCompileRequest.of(selectSql(6), scope(fourSourcePolicy)));
|
||||||
|
Assert.fail("expected source limit rejection");
|
||||||
|
} catch (FederationSqlException exception) {
|
||||||
|
Assert.assertEquals(FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
|
||||||
|
exception.errorCode());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private FederationQueryScopeDefinition scope(FederationExecutionPolicy policy) {
|
||||||
|
return new FederationQueryScopeDefinition(
|
||||||
|
"six-source-scope", 1, bindings, "S1", policy);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String selectSql(int sourceCount) {
|
||||||
|
StringBuilder sql = new StringBuilder("SELECT s1.SUBJECT_ID");
|
||||||
|
for (int index = 1; index <= sourceCount; index++) {
|
||||||
|
sql.append(", s").append(index).append(".METRIC_VALUE");
|
||||||
|
}
|
||||||
|
sql.append(" FROM S1.APP.SUBJECT_METRIC s1");
|
||||||
|
for (int index = 2; index <= sourceCount; index++) {
|
||||||
|
sql.append(" JOIN S").append(index).append(".APP.SUBJECT_METRIC s")
|
||||||
|
.append(index).append(" ON s").append(index)
|
||||||
|
.append(".SUBJECT_ID = s1.SUBJECT_ID");
|
||||||
|
}
|
||||||
|
return sql.append(" ORDER BY s1.SUBJECT_ID").toString();
|
||||||
|
}
|
||||||
|
|
||||||
|
private static JdbcDataSource dataSource(String name) {
|
||||||
|
JdbcDataSource dataSource = new JdbcDataSource();
|
||||||
|
dataSource.setURL("jdbc:h2:mem:" + name + ";DB_CLOSE_DELAY=-1");
|
||||||
|
dataSource.setUser("sa");
|
||||||
|
dataSource.setPassword("");
|
||||||
|
return dataSource;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static void execute(JdbcDataSource dataSource, String... statements)
|
||||||
|
throws Exception {
|
||||||
|
try (Connection connection = dataSource.getConnection();
|
||||||
|
Statement statement = connection.createStatement()) {
|
||||||
|
for (String sql : statements) {
|
||||||
|
statement.execute(sql);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static FederationSourceDefinition definition(SourceId sourceId) {
|
||||||
|
return new FederationSourceDefinition(
|
||||||
|
sourceId, 1, JdbcFederationSqlAdapterProvider.ADAPTER_ID,
|
||||||
|
List.of(new JdbcSchemaDefinition("APP", null, "PUBLIC")), Map.of());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -70,6 +70,7 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
|
|||||||
private final CalciteFederationSqlCompiler compiler;
|
private final CalciteFederationSqlCompiler compiler;
|
||||||
private final AdapterFederationStatisticsProvider automaticStatisticsProvider;
|
private final AdapterFederationStatisticsProvider automaticStatisticsProvider;
|
||||||
private final NodeMemoryAdmissionController nodeMemoryAdmission;
|
private final NodeMemoryAdmissionController nodeMemoryAdmission;
|
||||||
|
private final FederationFragmentScheduler fragmentScheduler;
|
||||||
private final CalciteSqlCompleter completer = new CalciteSqlCompleter();
|
private final CalciteSqlCompleter completer = new CalciteSqlCompleter();
|
||||||
private final QueryCancellationRegistry cancellations = new QueryCancellationRegistry();
|
private final QueryCancellationRegistry cancellations = new QueryCancellationRegistry();
|
||||||
private final ScheduledThreadPoolExecutor deadlineScheduler = deadlineScheduler();
|
private final ScheduledThreadPoolExecutor deadlineScheduler = deadlineScheduler();
|
||||||
@@ -297,6 +298,7 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
|
|||||||
this.nodeMemoryAdmission = new NodeMemoryAdmissionController(
|
this.nodeMemoryAdmission = new NodeMemoryAdmissionController(
|
||||||
maximumNodeIntermediateBytes
|
maximumNodeIntermediateBytes
|
||||||
);
|
);
|
||||||
|
this.fragmentScheduler = new FederationFragmentScheduler(executionPolicy);
|
||||||
this.planCache = new BoundedPlanCache(
|
this.planCache = new BoundedPlanCache(
|
||||||
maximumPlanCacheEntries,
|
maximumPlanCacheEntries,
|
||||||
maximumConcurrentCompilations,
|
maximumConcurrentCompilations,
|
||||||
@@ -650,7 +652,8 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
|
|||||||
registration,
|
registration,
|
||||||
compiler.effectivePolicy(plan.queryScope()),
|
compiler.effectivePolicy(plan.queryScope()),
|
||||||
metrics,
|
metrics,
|
||||||
queryDeadline
|
queryDeadline,
|
||||||
|
fragmentScheduler
|
||||||
);
|
);
|
||||||
queryCursor = new FederatedResultCursor(
|
queryCursor = new FederatedResultCursor(
|
||||||
plan,
|
plan,
|
||||||
@@ -1080,6 +1083,11 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
|
|||||||
} catch (RuntimeException exception) {
|
} catch (RuntimeException exception) {
|
||||||
failure = exception;
|
failure = exception;
|
||||||
}
|
}
|
||||||
|
try {
|
||||||
|
fragmentScheduler.close();
|
||||||
|
} catch (RuntimeException exception) {
|
||||||
|
failure = append(failure, exception);
|
||||||
|
}
|
||||||
try {
|
try {
|
||||||
sourceManager.close();
|
sourceManager.close();
|
||||||
} catch (RuntimeException exception) {
|
} catch (RuntimeException exception) {
|
||||||
|
|||||||
@@ -62,6 +62,7 @@ final class FederatedResultCursor
|
|||||||
this.session = session;
|
this.session = session;
|
||||||
this.metrics = metrics;
|
this.metrics = metrics;
|
||||||
this.maximumRows = maximumRows;
|
this.maximumRows = maximumRows;
|
||||||
|
session.prepareFragments();
|
||||||
this.enumerator = plan.localBindable()
|
this.enumerator = plan.localBindable()
|
||||||
.bind(new FederationDataContext(
|
.bind(new FederationDataContext(
|
||||||
plan.localRootSchema(),
|
plan.localRootSchema(),
|
||||||
|
|||||||
@@ -20,14 +20,21 @@ import java.time.LocalDateTime;
|
|||||||
import java.time.LocalTime;
|
import java.time.LocalTime;
|
||||||
import java.time.OffsetDateTime;
|
import java.time.OffsetDateTime;
|
||||||
import java.time.ZoneOffset;
|
import java.time.ZoneOffset;
|
||||||
|
import java.util.ArrayDeque;
|
||||||
import java.util.ArrayList;
|
import java.util.ArrayList;
|
||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.Map;
|
import java.util.Map;
|
||||||
import java.util.Set;
|
import java.util.Set;
|
||||||
|
import java.util.concurrent.CompletableFuture;
|
||||||
import java.util.concurrent.ConcurrentHashMap;
|
import java.util.concurrent.ConcurrentHashMap;
|
||||||
|
import java.util.concurrent.ExecutionException;
|
||||||
|
import java.util.concurrent.Future;
|
||||||
|
import java.util.concurrent.FutureTask;
|
||||||
import java.util.concurrent.Semaphore;
|
import java.util.concurrent.Semaphore;
|
||||||
import java.util.concurrent.TimeUnit;
|
import java.util.concurrent.TimeUnit;
|
||||||
|
import java.util.concurrent.TimeoutException;
|
||||||
import java.util.concurrent.atomic.AtomicBoolean;
|
import java.util.concurrent.atomic.AtomicBoolean;
|
||||||
|
import java.util.concurrent.atomic.AtomicReference;
|
||||||
import org.apache.calcite.avatica.util.ByteString;
|
import org.apache.calcite.avatica.util.ByteString;
|
||||||
import org.apache.calcite.linq4j.AbstractEnumerable;
|
import org.apache.calcite.linq4j.AbstractEnumerable;
|
||||||
import org.apache.calcite.linq4j.Enumerable;
|
import org.apache.calcite.linq4j.Enumerable;
|
||||||
@@ -50,9 +57,19 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
private final FederationExecutionPolicy policy;
|
private final FederationExecutionPolicy policy;
|
||||||
private final FederationQueryMetricsTracker metrics;
|
private final FederationQueryMetricsTracker metrics;
|
||||||
private final QueryDeadline deadline;
|
private final QueryDeadline deadline;
|
||||||
|
private final FederationFragmentScheduler fragmentScheduler;
|
||||||
private final Semaphore fragmentSlots;
|
private final Semaphore fragmentSlots;
|
||||||
private final Set<FederationResultCursor> openFragmentCursors = ConcurrentHashMap.newKeySet();
|
private final Set<FederationResultCursor> openFragmentCursors = ConcurrentHashMap.newKeySet();
|
||||||
|
private final Set<Future<?>> fragmentTasks = ConcurrentHashMap.newKeySet();
|
||||||
|
private final Map<String, CompletableFuture<PreparedFragment>> fragmentPreparations =
|
||||||
|
new ConcurrentHashMap<>();
|
||||||
|
private final Object preparationLock = new Object();
|
||||||
|
private final ArrayDeque<CompiledFederationFragment> pendingPreparations =
|
||||||
|
new ArrayDeque<>();
|
||||||
private final AtomicBoolean closed = new AtomicBoolean();
|
private final AtomicBoolean closed = new AtomicBoolean();
|
||||||
|
private final AtomicReference<Throwable> preparationFailure = new AtomicReference<>();
|
||||||
|
private int runningPreparations;
|
||||||
|
private boolean preparationFailed;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 创建查询执行会话。
|
* 创建查询执行会话。
|
||||||
@@ -64,6 +81,7 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
* @param policy 有效资源策略
|
* @param policy 有效资源策略
|
||||||
* @param metrics 指标跟踪器
|
* @param metrics 指标跟踪器
|
||||||
* @param deadline 请求级统一截止时间
|
* @param deadline 请求级统一截止时间
|
||||||
|
* @param fragmentScheduler Engine 级有界分片调度器
|
||||||
*/
|
*/
|
||||||
FederationExecutionSession(
|
FederationExecutionSession(
|
||||||
DefaultFederationSqlPlan plan,
|
DefaultFederationSqlPlan plan,
|
||||||
@@ -72,7 +90,8 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
QueryCancellationRegistry.QueryRegistration registration,
|
QueryCancellationRegistry.QueryRegistration registration,
|
||||||
FederationExecutionPolicy policy,
|
FederationExecutionPolicy policy,
|
||||||
FederationQueryMetricsTracker metrics,
|
FederationQueryMetricsTracker metrics,
|
||||||
QueryDeadline deadline
|
QueryDeadline deadline,
|
||||||
|
FederationFragmentScheduler fragmentScheduler
|
||||||
) {
|
) {
|
||||||
this.plan = plan;
|
this.plan = plan;
|
||||||
this.context = context;
|
this.context = context;
|
||||||
@@ -81,9 +100,41 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
this.policy = policy;
|
this.policy = policy;
|
||||||
this.metrics = metrics;
|
this.metrics = metrics;
|
||||||
this.deadline = deadline;
|
this.deadline = deadline;
|
||||||
|
this.fragmentScheduler = fragmentScheduler;
|
||||||
this.fragmentSlots = new Semaphore(policy.maximumConcurrentFragments(), true);
|
this.fragmentSlots = new Semaphore(policy.maximumConcurrentFragments(), true);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 异步启动多源物理分片,在查询级中间结果预算内完成有界物化。
|
||||||
|
*/
|
||||||
|
void prepareFragments() {
|
||||||
|
if (plan.fragments().size() < 2 || policy.maximumConcurrentFragments() < 2) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
for (FederationFragmentPlan fragment : plan.fragments()) {
|
||||||
|
CompiledFederationFragment compiled = plan.compiledFragments().get(
|
||||||
|
fragment.fragmentId()
|
||||||
|
);
|
||||||
|
if (compiled == null) {
|
||||||
|
throw new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
|
"compiled fragment is missing: " + fragment.fragmentId()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
CompletableFuture<PreparedFragment> preparation = new CompletableFuture<>();
|
||||||
|
fragmentPreparations.put(fragment.fragmentId(), preparation);
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
pendingPreparations.addLast(compiled);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
schedulePreparations();
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
cancelPreparation();
|
||||||
|
throw exception;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 为 Calcite ScannableTable 创建一个按需执行物理分片的 Enumerable。
|
* 为 Calcite ScannableTable 创建一个按需执行物理分片的 Enumerable。
|
||||||
*
|
*
|
||||||
@@ -101,6 +152,13 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
return new AbstractEnumerable<>() {
|
return new AbstractEnumerable<>() {
|
||||||
@Override
|
@Override
|
||||||
public Enumerator<Object[]> enumerator() {
|
public Enumerator<Object[]> enumerator() {
|
||||||
|
CompletableFuture<PreparedFragment> preparation =
|
||||||
|
fragmentPreparations.get(fragmentId);
|
||||||
|
if (preparation != null) {
|
||||||
|
PreparedFragment prepared = awaitPrepared(preparation);
|
||||||
|
fragmentPreparations.remove(fragmentId, preparation);
|
||||||
|
return prepared.enumerator();
|
||||||
|
}
|
||||||
return openFragment(compiled);
|
return openFragment(compiled);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -138,6 +196,7 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
}
|
}
|
||||||
registration.ensureNotCancelled();
|
registration.ensureNotCancelled();
|
||||||
}
|
}
|
||||||
|
throwPreparationFailure();
|
||||||
deadline.ensureAllowed();
|
deadline.ensureAllowed();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -166,9 +225,189 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
private Enumerator<Object[]> openFragment(CompiledFederationFragment compiled) {
|
private Enumerator<Object[]> openFragment(CompiledFederationFragment compiled) {
|
||||||
ensureExecutionAllowed();
|
ensureExecutionAllowed();
|
||||||
acquireFragmentSlot();
|
acquireFragmentSlot();
|
||||||
|
try {
|
||||||
|
return new FragmentEnumerator(compiled, createFragmentCursor(compiled), true);
|
||||||
|
} catch (RuntimeException exception) {
|
||||||
|
fragmentSlots.release();
|
||||||
|
throw exception;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void schedulePreparations() {
|
||||||
|
List<PreparationTask> tasks = new ArrayList<>();
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
while (!closed.get()
|
||||||
|
&& !preparationFailed
|
||||||
|
&& runningPreparations < policy.maximumConcurrentFragments()
|
||||||
|
&& !pendingPreparations.isEmpty()) {
|
||||||
|
CompiledFederationFragment compiled = pendingPreparations.removeFirst();
|
||||||
|
CompletableFuture<PreparedFragment> preparation =
|
||||||
|
fragmentPreparations.get(
|
||||||
|
compiled.plan().fragmentId()
|
||||||
|
);
|
||||||
|
if (preparation == null || preparation.isCancelled()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
runningPreparations++;
|
||||||
|
tasks.add(new PreparationTask(compiled, preparation));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (PreparationTask task : tasks) {
|
||||||
|
submitPreparation(task);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void submitPreparation(PreparationTask preparationTask) {
|
||||||
|
FutureTask<Void> task = new FutureTask<>(() -> {
|
||||||
|
runPreparation(preparationTask);
|
||||||
|
return null;
|
||||||
|
}) {
|
||||||
|
@Override
|
||||||
|
protected void done() {
|
||||||
|
fragmentTasks.remove(this);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
fragmentTasks.add(task);
|
||||||
|
try {
|
||||||
|
fragmentScheduler.execute(task);
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
fragmentTasks.remove(task);
|
||||||
|
task.cancel(true);
|
||||||
|
preparationTask.preparation().completeExceptionally(exception);
|
||||||
|
failPreparations(exception);
|
||||||
|
preparationFinished();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void runPreparation(PreparationTask task) {
|
||||||
|
try {
|
||||||
|
task.preparation().complete(prepareFragment(task.compiled()));
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
task.preparation().completeExceptionally(exception);
|
||||||
|
failPreparations(exception);
|
||||||
|
} finally {
|
||||||
|
preparationFinished();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void preparationFinished() {
|
||||||
|
boolean scheduleNext;
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
runningPreparations--;
|
||||||
|
scheduleNext = !closed.get() && !preparationFailed;
|
||||||
|
}
|
||||||
|
if (scheduleNext) {
|
||||||
|
schedulePreparations();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void failPreparations(Throwable failure) {
|
||||||
|
if (!preparationFailure.compareAndSet(null, failure)) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
preparationFailed = true;
|
||||||
|
pendingPreparations.clear();
|
||||||
|
}
|
||||||
|
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
|
||||||
|
preparation.completeExceptionally(failure);
|
||||||
|
}
|
||||||
|
registration.terminateStatementsForCleanup();
|
||||||
|
}
|
||||||
|
|
||||||
|
private void throwPreparationFailure() {
|
||||||
|
Throwable failure = preparationFailure.get();
|
||||||
|
if (failure == null) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (failure instanceof RuntimeException runtimeException) {
|
||||||
|
throw runtimeException;
|
||||||
|
}
|
||||||
|
if (failure instanceof Error error) {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
throw new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
|
"federation fragment preparation failed",
|
||||||
|
failure
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
private PreparedFragment awaitPrepared(Future<PreparedFragment> preparation) {
|
||||||
|
try {
|
||||||
|
while (true) {
|
||||||
|
ensureExecutionAllowed();
|
||||||
|
long remaining = deadline.remainingNanos();
|
||||||
|
if (remaining <= 0L) {
|
||||||
|
ensureExecutionAllowed();
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
return preparation.get(
|
||||||
|
Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(50)),
|
||||||
|
TimeUnit.NANOSECONDS
|
||||||
|
);
|
||||||
|
} catch (TimeoutException ignored) {
|
||||||
|
// 周期性复核共享 Deadline 和取消状态。
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} catch (InterruptedException exception) {
|
||||||
|
Thread.currentThread().interrupt();
|
||||||
|
cancelPreparation();
|
||||||
|
throw new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
|
"waiting for federation fragment preparation was interrupted",
|
||||||
|
exception
|
||||||
|
);
|
||||||
|
} catch (ExecutionException exception) {
|
||||||
|
cancelPreparation();
|
||||||
|
Throwable cause = exception.getCause();
|
||||||
|
if (cause instanceof RuntimeException runtimeException) {
|
||||||
|
throw runtimeException;
|
||||||
|
}
|
||||||
|
if (cause instanceof Error error) {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
throw new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
|
"federation fragment preparation failed",
|
||||||
|
cause
|
||||||
|
);
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
cancelPreparation();
|
||||||
|
throw exception;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private PreparedFragment prepareFragment(CompiledFederationFragment compiled) {
|
||||||
|
ensureExecutionAllowed();
|
||||||
|
FederationResultCursor cursor = null;
|
||||||
|
try {
|
||||||
|
cursor = createFragmentCursor(compiled);
|
||||||
|
List<Object[]> rows = new ArrayList<>();
|
||||||
|
while (true) {
|
||||||
|
ensureExecutionAllowed();
|
||||||
|
if (!cursor.next()) {
|
||||||
|
releaseFragment(compiled, cursor, true, false);
|
||||||
|
cursor = null;
|
||||||
|
return new PreparedFragment(rows);
|
||||||
|
}
|
||||||
|
Object[] row = normalizeRow(cursor, compiled.rowType());
|
||||||
|
recordFragmentIntermediate(compiled, row);
|
||||||
|
rows.add(row);
|
||||||
|
}
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
if (cursor != null) {
|
||||||
|
releaseFragment(compiled, cursor, false, false);
|
||||||
|
}
|
||||||
|
throw exception;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private FederationResultCursor createFragmentCursor(
|
||||||
|
CompiledFederationFragment compiled
|
||||||
|
) {
|
||||||
FederationFragmentPlan fragment = compiled.plan();
|
FederationFragmentPlan fragment = compiled.plan();
|
||||||
SourceRuntime runtime = runtimeSnapshot.runtime(fragment.bindingName());
|
SourceRuntime runtime = runtimeSnapshot.runtime(fragment.bindingName());
|
||||||
try {
|
|
||||||
FederationResultCursor cursor = runtime.adapter().fragmentExecutor().execute(
|
FederationResultCursor cursor = runtime.adapter().fragmentExecutor().execute(
|
||||||
new FederationFragmentExecutionContext(
|
new FederationFragmentExecutionContext(
|
||||||
context.queryId(),
|
context.queryId(),
|
||||||
@@ -191,13 +430,55 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
openFragmentCursors.add(cursor);
|
openFragmentCursors.add(cursor);
|
||||||
return new FragmentEnumerator(compiled, cursor);
|
try {
|
||||||
} catch (RuntimeException exception) {
|
ensureExecutionAllowed();
|
||||||
fragmentSlots.release();
|
return cursor;
|
||||||
|
} catch (RuntimeException | Error exception) {
|
||||||
|
openFragmentCursors.remove(cursor);
|
||||||
|
closeQuietly(cursor);
|
||||||
throw exception;
|
throw exception;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private void recordFragmentIntermediate(
|
||||||
|
CompiledFederationFragment compiled,
|
||||||
|
Object[] row
|
||||||
|
) {
|
||||||
|
FederationQueryMetricsTracker.IntermediateUsage usage =
|
||||||
|
metrics.recordIntermediate(compiled.plan().fragmentId(), row);
|
||||||
|
if (usage.rows() > policy.maximumIntermediateRows()
|
||||||
|
|| usage.bytes() > policy.maximumIntermediateBytes()) {
|
||||||
|
throw new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
|
||||||
|
"federation intermediate result limit was exceeded"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private void cancelPreparation() {
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
preparationFailed = true;
|
||||||
|
pendingPreparations.clear();
|
||||||
|
}
|
||||||
|
for (Future<?> task : List.copyOf(fragmentTasks)) {
|
||||||
|
task.cancel(true);
|
||||||
|
}
|
||||||
|
fragmentTasks.clear();
|
||||||
|
for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
|
||||||
|
openFragmentCursors.remove(cursor);
|
||||||
|
try {
|
||||||
|
cursor.close();
|
||||||
|
} catch (RuntimeException ignored) {
|
||||||
|
// 保留触发取消的原始分片异常;Statement 取消通道继续收口剩余资源。
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
|
||||||
|
preparation.cancel(true);
|
||||||
|
}
|
||||||
|
fragmentPreparations.clear();
|
||||||
|
registration.terminateStatementsForCleanup();
|
||||||
|
}
|
||||||
|
|
||||||
private void acquireFragmentSlot() {
|
private void acquireFragmentSlot() {
|
||||||
while (true) {
|
while (true) {
|
||||||
ensureExecutionAllowed();
|
ensureExecutionAllowed();
|
||||||
@@ -259,6 +540,21 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
if (!closed.compareAndSet(false, true)) {
|
if (!closed.compareAndSet(false, true)) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
boolean terminationRequired = !fragmentTasks.isEmpty()
|
||||||
|
|| !openFragmentCursors.isEmpty();
|
||||||
|
synchronized (preparationLock) {
|
||||||
|
terminationRequired = terminationRequired || runningPreparations > 0;
|
||||||
|
preparationFailed = true;
|
||||||
|
pendingPreparations.clear();
|
||||||
|
}
|
||||||
|
for (Future<?> task : List.copyOf(fragmentTasks)) {
|
||||||
|
task.cancel(true);
|
||||||
|
}
|
||||||
|
fragmentTasks.clear();
|
||||||
|
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
|
||||||
|
preparation.cancel(true);
|
||||||
|
}
|
||||||
|
fragmentPreparations.clear();
|
||||||
RuntimeException failure = null;
|
RuntimeException failure = null;
|
||||||
for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
|
for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
|
||||||
try {
|
try {
|
||||||
@@ -273,6 +569,9 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
openFragmentCursors.remove(cursor);
|
openFragmentCursors.remove(cursor);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if (terminationRequired) {
|
||||||
|
registration.terminateStatementsForCleanup();
|
||||||
|
}
|
||||||
if (failure != null) {
|
if (failure != null) {
|
||||||
throw failure;
|
throw failure;
|
||||||
}
|
}
|
||||||
@@ -282,15 +581,18 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
|
|
||||||
private final CompiledFederationFragment compiled;
|
private final CompiledFederationFragment compiled;
|
||||||
private final FederationResultCursor cursor;
|
private final FederationResultCursor cursor;
|
||||||
|
private final boolean fragmentSlotHeld;
|
||||||
private final AtomicBoolean released = new AtomicBoolean();
|
private final AtomicBoolean released = new AtomicBoolean();
|
||||||
private Object[] current;
|
private Object[] current;
|
||||||
|
|
||||||
private FragmentEnumerator(
|
private FragmentEnumerator(
|
||||||
CompiledFederationFragment compiled,
|
CompiledFederationFragment compiled,
|
||||||
FederationResultCursor cursor
|
FederationResultCursor cursor,
|
||||||
|
boolean fragmentSlotHeld
|
||||||
) {
|
) {
|
||||||
this.compiled = compiled;
|
this.compiled = compiled;
|
||||||
this.cursor = cursor;
|
this.cursor = cursor;
|
||||||
|
this.fragmentSlotHeld = fragmentSlotHeld;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** {@inheritDoc} */
|
/** {@inheritDoc} */
|
||||||
@@ -309,15 +611,7 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
Object[] row = normalizeRow(cursor, compiled.rowType());
|
Object[] row = normalizeRow(cursor, compiled.rowType());
|
||||||
FederationQueryMetricsTracker.IntermediateUsage usage =
|
recordFragmentIntermediate(compiled, row);
|
||||||
metrics.recordIntermediate(compiled.plan().fragmentId(), row);
|
|
||||||
if (usage.rows() > policy.maximumIntermediateRows()
|
|
||||||
|| usage.bytes() > policy.maximumIntermediateBytes()) {
|
|
||||||
throw new FederationSqlException(
|
|
||||||
FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
|
|
||||||
"federation intermediate result limit was exceeded"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
current = row;
|
current = row;
|
||||||
return true;
|
return true;
|
||||||
} catch (RuntimeException exception) {
|
} catch (RuntimeException exception) {
|
||||||
@@ -346,7 +640,9 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
try {
|
try {
|
||||||
cursor.close();
|
cursor.close();
|
||||||
} finally {
|
} finally {
|
||||||
|
if (fragmentSlotHeld) {
|
||||||
fragmentSlots.release();
|
fragmentSlots.release();
|
||||||
|
}
|
||||||
if (exhausted) {
|
if (exhausted) {
|
||||||
metrics.finishFragment(compiled.plan().fragmentId());
|
metrics.finishFragment(compiled.plan().fragmentId());
|
||||||
}
|
}
|
||||||
@@ -354,6 +650,78 @@ final class FederationExecutionSession implements AutoCloseable {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private static final class PreparedFragment {
|
||||||
|
|
||||||
|
private final List<Object[]> rows;
|
||||||
|
|
||||||
|
private PreparedFragment(List<Object[]> rows) {
|
||||||
|
this.rows = rows;
|
||||||
|
}
|
||||||
|
|
||||||
|
private Enumerator<Object[]> enumerator() {
|
||||||
|
return new Enumerator<>() {
|
||||||
|
|
||||||
|
private int index = -1;
|
||||||
|
private Object[] current;
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Object[] current() {
|
||||||
|
return current;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public boolean moveNext() {
|
||||||
|
int next = index + 1;
|
||||||
|
if (next >= rows.size()) {
|
||||||
|
current = null;
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
index = next;
|
||||||
|
current = rows.set(index, null);
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void reset() {
|
||||||
|
throw new UnsupportedOperationException(
|
||||||
|
"fragment cursor cannot be reset"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void close() {
|
||||||
|
rows.clear();
|
||||||
|
current = null;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private record PreparationTask(
|
||||||
|
CompiledFederationFragment compiled,
|
||||||
|
CompletableFuture<PreparedFragment> preparation
|
||||||
|
) {
|
||||||
|
}
|
||||||
|
|
||||||
|
private void releaseFragment(
|
||||||
|
CompiledFederationFragment compiled,
|
||||||
|
FederationResultCursor cursor,
|
||||||
|
boolean exhausted,
|
||||||
|
boolean fragmentSlotHeld
|
||||||
|
) {
|
||||||
|
openFragmentCursors.remove(cursor);
|
||||||
|
try {
|
||||||
|
cursor.close();
|
||||||
|
} finally {
|
||||||
|
if (fragmentSlotHeld) {
|
||||||
|
fragmentSlots.release();
|
||||||
|
}
|
||||||
|
if (exhausted) {
|
||||||
|
metrics.finishFragment(compiled.plan().fragmentId());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
private static Object[] normalizeRow(
|
private static Object[] normalizeRow(
|
||||||
FederationResultCursor cursor,
|
FederationResultCursor cursor,
|
||||||
RelDataType rowType
|
RelDataType rowType
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
package com.easyagents.federation.sql.runtime;
|
||||||
|
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlException;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
|
||||||
|
import java.util.concurrent.ArrayBlockingQueue;
|
||||||
|
import java.util.concurrent.Executor;
|
||||||
|
import java.util.concurrent.RejectedExecutionException;
|
||||||
|
import java.util.concurrent.ThreadPoolExecutor;
|
||||||
|
import java.util.concurrent.TimeUnit;
|
||||||
|
import java.util.concurrent.atomic.AtomicBoolean;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Engine 级有界分片工作线程池;查询级并发度仍由有效执行策略单独控制。
|
||||||
|
*/
|
||||||
|
final class FederationFragmentScheduler implements Executor, AutoCloseable {
|
||||||
|
|
||||||
|
private static final int MAXIMUM_WORKERS = 32;
|
||||||
|
private static final int MINIMUM_QUEUE_CAPACITY = 64;
|
||||||
|
private static final int QUEUE_CAPACITY_PER_WORKER = 8;
|
||||||
|
|
||||||
|
private final ThreadPoolExecutor executor;
|
||||||
|
private final AtomicBoolean closed = new AtomicBoolean();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 创建 Engine 独占的有界分片调度器。
|
||||||
|
*
|
||||||
|
* @param policy Engine 联邦执行硬上限
|
||||||
|
*/
|
||||||
|
FederationFragmentScheduler(FederationExecutionPolicy policy) {
|
||||||
|
int availableProcessors = Math.max(1, Runtime.getRuntime().availableProcessors());
|
||||||
|
int workers = Math.min(
|
||||||
|
MAXIMUM_WORKERS,
|
||||||
|
Math.max(policy.maximumConcurrentFragments(), availableProcessors)
|
||||||
|
);
|
||||||
|
int queueCapacity = Math.max(
|
||||||
|
MINIMUM_QUEUE_CAPACITY,
|
||||||
|
workers * QUEUE_CAPACITY_PER_WORKER
|
||||||
|
);
|
||||||
|
this.executor = executor(workers, queueCapacity);
|
||||||
|
}
|
||||||
|
|
||||||
|
FederationFragmentScheduler(int workers, int queueCapacity) {
|
||||||
|
if (workers <= 0) {
|
||||||
|
throw new IllegalArgumentException("workers must be positive");
|
||||||
|
}
|
||||||
|
if (queueCapacity <= 0) {
|
||||||
|
throw new IllegalArgumentException("queueCapacity must be positive");
|
||||||
|
}
|
||||||
|
this.executor = executor(workers, queueCapacity);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ThreadPoolExecutor executor(int workers, int queueCapacity) {
|
||||||
|
AtomicInteger threadSequence = new AtomicInteger();
|
||||||
|
return new ThreadPoolExecutor(
|
||||||
|
workers,
|
||||||
|
workers,
|
||||||
|
0L,
|
||||||
|
TimeUnit.MILLISECONDS,
|
||||||
|
new ArrayBlockingQueue<>(queueCapacity),
|
||||||
|
runnable -> {
|
||||||
|
Thread thread = new Thread(
|
||||||
|
runnable,
|
||||||
|
"easy-agents-federation-fragment-" + threadSequence.incrementAndGet()
|
||||||
|
);
|
||||||
|
thread.setDaemon(true);
|
||||||
|
return thread;
|
||||||
|
},
|
||||||
|
new ThreadPoolExecutor.AbortPolicy()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 提交分片任务;队列饱和时明确拒绝,禁止在调用线程绕过 Worker 上限执行 JDBC。
|
||||||
|
*
|
||||||
|
* @param command 分片任务
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void execute(Runnable command) {
|
||||||
|
if (closed.get()) {
|
||||||
|
throw engineClosed();
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
executor.execute(command);
|
||||||
|
} catch (RejectedExecutionException exception) {
|
||||||
|
if (closed.get() || executor.isShutdown()) {
|
||||||
|
throw engineClosed();
|
||||||
|
}
|
||||||
|
throw schedulerOverloaded(exception);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 中断排队任务并关闭工作线程。
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void close() {
|
||||||
|
if (closed.compareAndSet(false, true)) {
|
||||||
|
executor.shutdownNow();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static FederationSqlException engineClosed() {
|
||||||
|
return new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.ENGINE_CLOSED,
|
||||||
|
"federation fragment scheduler is closed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static FederationSqlException schedulerOverloaded(Throwable cause) {
|
||||||
|
return new FederationSqlException(
|
||||||
|
FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
|
||||||
|
"federation fragment scheduler is saturated",
|
||||||
|
cause
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -301,6 +301,13 @@ final class QueryCancellationRegistry implements AutoCloseable {
|
|||||||
return state.cancellationRequested();
|
return state.cancellationRequested();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 不改变查询终态,仅终止当前已登记的 Statement,用于提前关闭异步分片。
|
||||||
|
*/
|
||||||
|
void terminateStatementsForCleanup() {
|
||||||
|
state.terminateStatements();
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 已取消时立即抛出稳定错误。
|
* 已取消时立即抛出稳定错误。
|
||||||
*/
|
*/
|
||||||
@@ -485,20 +492,26 @@ final class QueryCancellationRegistry implements AutoCloseable {
|
|||||||
}
|
}
|
||||||
|
|
||||||
private void terminate(QueryId queryId) {
|
private void terminate(QueryId queryId) {
|
||||||
|
if (isClosed()) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
SQLException failure = null;
|
SQLException failure = null;
|
||||||
if (cancelIssued.compareAndSet(false, true)) {
|
if (cancelIssued.compareAndSet(false, true)) {
|
||||||
try {
|
try {
|
||||||
statement.cancel();
|
statement.cancel();
|
||||||
} catch (SQLException exception) {
|
} catch (SQLException exception) {
|
||||||
|
if (!isClosed()) {
|
||||||
cancelIssued.set(false);
|
cancelIssued.set(false);
|
||||||
failure = exception;
|
failure = exception;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
// close() 同时覆盖 register 与 executeQuery 之间的 JDBC 取消空窗。
|
// close() 同时覆盖 register 与 executeQuery 之间的 JDBC 取消空窗。
|
||||||
if (closeIssued.compareAndSet(false, true)) {
|
if (closeIssued.compareAndSet(false, true)) {
|
||||||
try {
|
try {
|
||||||
statement.close();
|
statement.close();
|
||||||
} catch (SQLException exception) {
|
} catch (SQLException exception) {
|
||||||
|
if (!isClosed()) {
|
||||||
closeIssued.set(false);
|
closeIssued.set(false);
|
||||||
if (failure == null) {
|
if (failure == null) {
|
||||||
failure = exception;
|
failure = exception;
|
||||||
@@ -507,6 +520,7 @@ final class QueryCancellationRegistry implements AutoCloseable {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
if (failure != null) {
|
if (failure != null) {
|
||||||
throw new FederationSqlException(
|
throw new FederationSqlException(
|
||||||
FederationSqlErrorCode.EXECUTION_FAILED,
|
FederationSqlErrorCode.EXECUTION_FAILED,
|
||||||
@@ -515,6 +529,14 @@ final class QueryCancellationRegistry implements AutoCloseable {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private boolean isClosed() {
|
||||||
|
try {
|
||||||
|
return statement.isClosed();
|
||||||
|
} catch (SQLException ignored) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
private static FederationSqlException cancelled(QueryId queryId, Throwable cause) {
|
private static FederationSqlException cancelled(QueryId queryId, Throwable cause) {
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
package com.easyagents.federation.sql.runtime;
|
||||||
|
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
|
||||||
|
import com.easyagents.federation.sql.api.FederationSqlException;
|
||||||
|
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
|
||||||
|
import java.util.concurrent.CountDownLatch;
|
||||||
|
import java.util.concurrent.ExecutorCompletionService;
|
||||||
|
import java.util.concurrent.Future;
|
||||||
|
import java.util.concurrent.TimeUnit;
|
||||||
|
import java.util.concurrent.atomic.AtomicBoolean;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
import org.junit.Assert;
|
||||||
|
import org.junit.Test;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Engine 级有界分片调度器测试。
|
||||||
|
*/
|
||||||
|
public class FederationFragmentSchedulerTest {
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证两个分片任务能够真实重叠执行,且调度器关闭后拒绝新任务。
|
||||||
|
*
|
||||||
|
* @throws Exception 等待任务失败
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldExecuteFragmentTasksConcurrentlyAndRejectAfterClose() throws Exception {
|
||||||
|
FederationFragmentScheduler scheduler = new FederationFragmentScheduler(
|
||||||
|
FederationExecutionPolicy.basic()
|
||||||
|
);
|
||||||
|
CountDownLatch started = new CountDownLatch(2);
|
||||||
|
CountDownLatch release = new CountDownLatch(1);
|
||||||
|
AtomicInteger active = new AtomicInteger();
|
||||||
|
AtomicInteger maximumActive = new AtomicInteger();
|
||||||
|
ExecutorCompletionService<Void> completions = new ExecutorCompletionService<>(scheduler);
|
||||||
|
try {
|
||||||
|
Future<Void> first = completions.submit(
|
||||||
|
() -> runBlockingTask(started, release, active, maximumActive)
|
||||||
|
);
|
||||||
|
Future<Void> second = completions.submit(
|
||||||
|
() -> runBlockingTask(started, release, active, maximumActive)
|
||||||
|
);
|
||||||
|
|
||||||
|
Assert.assertTrue(started.await(2, TimeUnit.SECONDS));
|
||||||
|
Assert.assertEquals(2, maximumActive.get());
|
||||||
|
release.countDown();
|
||||||
|
first.get(2, TimeUnit.SECONDS);
|
||||||
|
second.get(2, TimeUnit.SECONDS);
|
||||||
|
} finally {
|
||||||
|
release.countDown();
|
||||||
|
scheduler.close();
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
scheduler.execute(() -> { });
|
||||||
|
Assert.fail("closed scheduler should reject new tasks");
|
||||||
|
} catch (FederationSqlException exception) {
|
||||||
|
Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, exception.errorCode());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 验证队列饱和时明确拒绝,且不会在提交线程绕过 Worker 上限执行任务。
|
||||||
|
*
|
||||||
|
* @throws Exception 等待任务失败
|
||||||
|
*/
|
||||||
|
@Test
|
||||||
|
public void shouldRejectSaturationWithoutRunningTaskOnCaller() throws Exception {
|
||||||
|
FederationFragmentScheduler scheduler = new FederationFragmentScheduler(1, 1);
|
||||||
|
CountDownLatch firstStarted = new CountDownLatch(1);
|
||||||
|
CountDownLatch releaseFirst = new CountDownLatch(1);
|
||||||
|
CountDownLatch queuedCompleted = new CountDownLatch(1);
|
||||||
|
AtomicBoolean rejectedTaskRan = new AtomicBoolean();
|
||||||
|
try {
|
||||||
|
scheduler.execute(() -> {
|
||||||
|
firstStarted.countDown();
|
||||||
|
try {
|
||||||
|
releaseFirst.await(2, TimeUnit.SECONDS);
|
||||||
|
} catch (InterruptedException exception) {
|
||||||
|
Thread.currentThread().interrupt();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Assert.assertTrue(firstStarted.await(2, TimeUnit.SECONDS));
|
||||||
|
scheduler.execute(queuedCompleted::countDown);
|
||||||
|
|
||||||
|
try {
|
||||||
|
scheduler.execute(() -> rejectedTaskRan.set(true));
|
||||||
|
Assert.fail("saturated scheduler should reject the task");
|
||||||
|
} catch (FederationSqlException exception) {
|
||||||
|
Assert.assertEquals(
|
||||||
|
FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
|
||||||
|
exception.errorCode()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Assert.assertFalse(rejectedTaskRan.get());
|
||||||
|
} finally {
|
||||||
|
releaseFirst.countDown();
|
||||||
|
Assert.assertTrue(queuedCompleted.await(2, TimeUnit.SECONDS));
|
||||||
|
scheduler.close();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static Void runBlockingTask(
|
||||||
|
CountDownLatch started,
|
||||||
|
CountDownLatch release,
|
||||||
|
AtomicInteger active,
|
||||||
|
AtomicInteger maximumActive
|
||||||
|
) throws InterruptedException {
|
||||||
|
int current = active.incrementAndGet();
|
||||||
|
maximumActive.accumulateAndGet(current, Math::max);
|
||||||
|
started.countDown();
|
||||||
|
try {
|
||||||
|
Assert.assertTrue(release.await(2, TimeUnit.SECONDS));
|
||||||
|
return null;
|
||||||
|
} finally {
|
||||||
|
active.decrementAndGet();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user