for #1172, refactor OpenTracingSQLExecutionHook

apache · Sep 25, 2018 · 3a3419e · 3a3419e
1 parent 9ded871
commit 3a3419e
Show file tree

Hide file tree

Showing 11 changed files with 31 additions and 28 deletions.
diff --git a/sharding-core/src/main/java/io/shardingsphere/core/executor/ShardingExecuteCallback.java b/sharding-core/src/main/java/io/shardingsphere/core/executor/ShardingExecuteCallback.java
@@ -33,8 +33,9 @@ public interface ShardingExecuteCallback<I, O> {
      * Execute callback.
      * 
      * @param input input value
+     * @param isTrunkThread is execution in trunk thread
      * @return execute result
      * @throws SQLException throw when execute failure
      */
-    O execute(I input) throws SQLException;
+    O execute(I input, boolean isTrunkThread) throws SQLException;
 }
diff --git a/sharding-core/src/main/java/io/shardingsphere/core/executor/ShardingExecuteEngine.java b/sharding-core/src/main/java/io/shardingsphere/core/executor/ShardingExecuteEngine.java
@@ -95,15 +95,15 @@ private <I, O> Collection<ListenableFuture<O>> asyncExecute(final Collection<I>
 
                 @Override
                 public O call() throws SQLException {
-                    return callback.execute(each);
+                    return callback.execute(each, false);
                 }
             }));
         }
         return result;
     }
 
     private <I, O> O syncExecute(final I input, final ShardingExecuteCallback<I, O> callback) throws SQLException {
-        return callback.execute(input);
+        return callback.execute(input, true);
     }
 
     private <O> List<O> getResults(final O firstResult, final Collection<ListenableFuture<O>> restFutures) throws SQLException {
@@ -168,13 +168,13 @@ private <I, O> ListenableFuture<Collection<O>> asyncGroupExecute(final ShardingE
 
             @Override
             public Collection<O> call() throws SQLException {
-                return callback.execute(inputGroup.getInputs());
+                return callback.execute(inputGroup.getInputs(), false);
             }
         });
     }
 
     private <I, O> Collection<O> syncGroupExecute(final ShardingExecuteGroup<I> executeGroup, final ShardingGroupExecuteCallback<I, O> callback) throws SQLException {
-        return callback.execute(executeGroup.getInputs());
+        return callback.execute(executeGroup.getInputs(), true);
     }
 
     private <O> List<O> getGroupResults(final Collection<O> firstResults, final Collection<ListenableFuture<Collection<O>>> restFutures) throws SQLException {

diff --git a/...ding-core/src/main/java/io/shardingsphere/core/executor/ShardingGroupExecuteCallback.java b/...ding-core/src/main/java/io/shardingsphere/core/executor/ShardingGroupExecuteCallback.java
@@ -34,8 +34,9 @@ public interface ShardingGroupExecuteCallback<I, O> {
      * Execute callback.
      * 
      * @param inputs input values
+     * @param isTrunkThread is execution in trunk thread
      * @return execute result
      * @throws SQLException throw when execute failure
      */
-    Collection<O> execute(Collection<I> inputs) throws SQLException;
+    Collection<O> execute(Collection<I> inputs, boolean isTrunkThread) throws SQLException;
 }
diff --git a/...ng-core/src/main/java/io/shardingsphere/core/executor/sql/execute/SQLExecuteCallback.java b/...ng-core/src/main/java/io/shardingsphere/core/executor/sql/execute/SQLExecuteCallback.java
@@ -64,26 +64,26 @@ public abstract class SQLExecuteCallback<T> implements ShardingExecuteCallback<S
     private final SQLExecutionHook sqlExecutionHook = new SPISQLExecutionHook();
 
     @Override
-    public final T execute(final StatementExecuteUnit statementExecuteUnit) throws SQLException {
-        return execute0(statementExecuteUnit);
+    public final T execute(final StatementExecuteUnit statementExecuteUnit, final boolean isTrunkThread) throws SQLException {
+        return execute0(statementExecuteUnit, isTrunkThread);
     }
 
     @Override
-    public final Collection<T> execute(final Collection<StatementExecuteUnit> statementExecuteUnits) throws SQLException {
+    public final Collection<T> execute(final Collection<StatementExecuteUnit> statementExecuteUnits, final boolean isTrunkThread) throws SQLException {
         Collection<T> result = new LinkedList<>();
         for (StatementExecuteUnit each : statementExecuteUnits) {
-            result.add(execute0(each));
+            result.add(execute0(each, isTrunkThread));
         }
         return result;
     }
 
-    private T execute0(final StatementExecuteUnit statementExecuteUnit) throws SQLException {
+    private T execute0(final StatementExecuteUnit statementExecuteUnit, final boolean isTrunkThread) throws SQLException {
         ExecutorExceptionHandler.setExceptionThrown(isExceptionThrown);
         ExecutorDataMap.setDataMap(dataMap);
         List<List<Object>> parameterSets = statementExecuteUnit.getRouteUnit().getSqlUnit().getParameterSets();
         DataSourceMetaData dataSourceMetaData = DataSourceMetaDataFactory.newInstance(databaseType, statementExecuteUnit.getStatement().getConnection().getMetaData().getURL());
         for (List<Object> each : parameterSets) {
-            sqlExecutionHook.start(statementExecuteUnit.getRouteUnit().getDataSourceName(), statementExecuteUnit.getRouteUnit().getSqlUnit().getSql(), each, dataSourceMetaData);
+            sqlExecutionHook.start(statementExecuteUnit.getRouteUnit().getDataSourceName(), statementExecuteUnit.getRouteUnit().getSqlUnit().getSql(), each, dataSourceMetaData, isTrunkThread);
             // TODO remove after BED removed
             shardingEventBus.post(SQLExecutionEventFactory.createEvent(sqlType, statementExecuteUnit, each, dataSourceMetaData));
         }

diff --git a/.../src/main/java/io/shardingsphere/core/executor/sql/prepare/SQLExecutePrepareTemplate.java b/.../src/main/java/io/shardingsphere/core/executor/sql/prepare/SQLExecutePrepareTemplate.java
@@ -81,7 +81,7 @@ private Collection<ShardingExecuteGroup<StatementExecuteUnit>> getAsynchronizedE
                 new ShardingExecuteCallback<Entry<String, List<SQLUnit>>, Collection<ShardingExecuteGroup<StatementExecuteUnit>>>() {
 
                     @Override
-                    public Collection<ShardingExecuteGroup<StatementExecuteUnit>> execute(final Entry<String, List<SQLUnit>> input) throws SQLException {
+                    public Collection<ShardingExecuteGroup<StatementExecuteUnit>> execute(final Entry<String, List<SQLUnit>> input, final boolean isTrunkThread) throws SQLException {
                         return getSQLExecuteGroups(input.getKey(), input.getValue(), callback);
                     }
                 });

diff --git a/...ore/src/main/java/io/shardingsphere/core/metadata/table/executor/TableMetaDataLoader.java b/...ore/src/main/java/io/shardingsphere/core/metadata/table/executor/TableMetaDataLoader.java
@@ -76,7 +76,7 @@ private List<TableMetaData> load(final Map<String, List<DataNode>> dataNodeGroup
         return executeEngine.groupExecute(getDataNodeGroups(dataNodeGroups), new ShardingGroupExecuteCallback<DataNode, TableMetaData>() {
 
             @Override
-            public Collection<TableMetaData> execute(final Collection<DataNode> dataNodes) throws SQLException {
+            public Collection<TableMetaData> execute(final Collection<DataNode> dataNodes, final boolean isTrunkThread) throws SQLException {
                 String dataSourceName = dataNodes.iterator().next().getDataSourceName();
                 DataSourceMetaData dataSourceMetaData = shardingDataSourceMetaData.getActualDataSourceMetaData(dataSourceName);
                 String catalog = null == dataSourceMetaData ? null : dataSourceMetaData.getSchemeName();

diff --git a/sharding-core/src/main/java/io/shardingsphere/core/spi/executor/SPISQLExecutionHook.java b/sharding-core/src/main/java/io/shardingsphere/core/spi/executor/SPISQLExecutionHook.java
@@ -32,9 +32,9 @@ public final class SPISQLExecutionHook implements SQLExecutionHook {
     private static final ServiceLoader<SQLExecutionHook> SERVICE_LOADER = ServiceLoader.load(SQLExecutionHook.class);
 
     @Override
-    public synchronized void start(final String dataSourceName, final String sql, final List<Object> parameters, final DataSourceMetaData dataSourceMetaData) {
+    public synchronized void start(final String dataSourceName, final String sql, final List<Object> parameters, final DataSourceMetaData dataSourceMetaData, final boolean isTrunkThread) {
         for (SQLExecutionHook each : SERVICE_LOADER) {
-            each.start(dataSourceName, sql, parameters, dataSourceMetaData);
+            each.start(dataSourceName, sql, parameters, dataSourceMetaData, isTrunkThread);
         }
     }
 

diff --git a/sharding-core/src/main/java/io/shardingsphere/core/spi/executor/SQLExecutionHook.java b/sharding-core/src/main/java/io/shardingsphere/core/spi/executor/SQLExecutionHook.java
@@ -35,8 +35,9 @@ public interface SQLExecutionHook {
      * @param sql SQL to be executed
      * @param parameters parameters of SQL
      * @param dataSourceMetaData data source meta data
+     * @param isTrunkThread is execution in trunk thread
      */
-    void start(String dataSourceName, String sql, List<Object> parameters, DataSourceMetaData dataSourceMetaData);
+    void start(String dataSourceName, String sql, List<Object> parameters, DataSourceMetaData dataSourceMetaData, boolean isTrunkThread);
 
     /**
      * Handle when SQL execution finished success.

diff --git a/...cing/src/main/java/io/shardingsphere/opentracing/handler/OpenTracingSQLExecutionHook.java b/...cing/src/main/java/io/shardingsphere/opentracing/handler/OpenTracingSQLExecutionHook.java
@@ -38,14 +38,14 @@ public final class OpenTracingSQLExecutionHook implements SQLExecutionHook {
 
     private static final String OPERATION_NAME = "/" + ShardingTags.COMPONENT_NAME + "/executeSQL/";
 
-    private final ThreadLocal<Boolean> isTrunkThread = new ThreadLocal<>();
+    private boolean isTrunkThread;
 
     private Span span;
 
     @Override
-    public void start(final String dataSourceName, final String sql, final List<Object> parameters, final DataSourceMetaData dataSourceMetaData) {
-        isTrunkThread.set(OpenTracingRootInvokeHandler.isTrunkThread());
-        if (ExecutorDataMap.getDataMap().containsKey(OpenTracingRootInvokeHandler.ROOT_SPAN_CONTINUATION) && !isTrunkThread.get()) {
+    public void start(final String dataSourceName, final String sql, final List<Object> parameters, final DataSourceMetaData dataSourceMetaData, final boolean isTrunkThread) {
+        this.isTrunkThread = isTrunkThread;
+        if (ExecutorDataMap.getDataMap().containsKey(OpenTracingRootInvokeHandler.ROOT_SPAN_CONTINUATION) && !isTrunkThread) {
             OpenTracingRootInvokeHandler.getActiveSpan().set(((ActiveSpan.Continuation) ExecutorDataMap.getDataMap().get(OpenTracingRootInvokeHandler.ROOT_SPAN_CONTINUATION)).activate());
         }
         span = ShardingTracer.get().buildSpan(OPERATION_NAME)
@@ -74,7 +74,7 @@ public void finishFailure(final Exception cause) {
     }
 
     private void deactivateSpan() {
-        if (!isTrunkThread.get()) {
+        if (!isTrunkThread) {
             OpenTracingRootInvokeHandler.getActiveSpan().get().deactivate();
             OpenTracingRootInvokeHandler.getActiveSpan().remove();
         }

diff --git a/...ding-opentracing/src/test/java/io/shardingsphere/opentracing/handler/AllHandlerTests.java b/...ding-opentracing/src/test/java/io/shardingsphere/opentracing/handler/AllHandlerTests.java
@@ -23,9 +23,9 @@
 
 @RunWith(Suite.class)
 @SuiteClasses({
-        OpenTracingRootInvokeHandlerTest.class,
-        OpenTracingParsingEventHandlerTest.class,
-        OpenTracingSQLExecutionEventHandlerTest.class
+        OpenTracingRootInvokeHandlerTest.class, 
+        OpenTracingParsingEventHandlerTest.class, 
+        OpenTracingSQLExecutionEventHandlerTest.class 
 })
 public final class AllHandlerTests {
 }
diff --git a/...t/java/io/shardingsphere/opentracing/handler/OpenTracingSQLExecutionEventHandlerTest.java b/...t/java/io/shardingsphere/opentracing/handler/OpenTracingSQLExecutionEventHandlerTest.java
@@ -51,7 +51,7 @@ public void assertExecuteSuccessForTrunkThread() {
         DataSourceMetaData dataSourceMetaData = mock(DataSourceMetaData.class);
         when(dataSourceMetaData.getHostName()).thenReturn("localhost");
         when(dataSourceMetaData.getPort()).thenReturn(8888);
-        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Arrays.<Object>asList("1", 2), dataSourceMetaData);
+        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Arrays.<Object>asList("1", 2), dataSourceMetaData, true);
         sqlExecutionHook.finishSuccess();
         assertThat(getTracer().finishedSpans().size(), is(1));
         MockSpan actual = getTracer().finishedSpans().get(0);
@@ -77,7 +77,7 @@ public void assertExecuteSuccessForBranchThread() {
         DataSourceMetaData dataSourceMetaData = mock(DataSourceMetaData.class);
         when(dataSourceMetaData.getHostName()).thenReturn("localhost");
         when(dataSourceMetaData.getPort()).thenReturn(8888);
-        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Arrays.<Object>asList("1", 2), dataSourceMetaData);
+        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Arrays.<Object>asList("1", 2), dataSourceMetaData, false);
         assertNotNull(OpenTracingRootInvokeHandler.getActiveSpan().get());
         sqlExecutionHook.finishSuccess();
         assertThat(getTracer().finishedSpans().size(), is(1));
@@ -102,7 +102,7 @@ public void assertExecuteFailure() {
         DataSourceMetaData dataSourceMetaData = mock(DataSourceMetaData.class);
         when(dataSourceMetaData.getHostName()).thenReturn("localhost");
         when(dataSourceMetaData.getPort()).thenReturn(8888);
-        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Collections.emptyList(), dataSourceMetaData);
+        sqlExecutionHook.start("ds_test", "SELECT * FROM XXX;", Collections.emptyList(), dataSourceMetaData, true);
         sqlExecutionHook.finishFailure(new RuntimeException("SQL execution error"));
         assertThat(getTracer().finishedSpans().size(), is(1));
         MockSpan actual = getTracer().finishedSpans().get(0);