Updated work with YdbTracer

Author: alex268

jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java

@@ -1,21 +1,18 @@
 package tech.ydb.jdbc;
 
 
-import tech.ydb.jdbc.impl.YdbTracerImpl;
-
-
 
 /**
  *
  * @author Aleksandr Gorshenin
  */
 public interface YdbTracer {
-    static void clear() {
-        YdbTracerImpl.clear();
-    }
+    interface Storage {
+        YdbTracer get();
+
+        default void clear() {
 
-    static YdbTracer current() {
-        return YdbTracerImpl.current();
+        }
     }
 
     void setId(String id);

jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java

@@ -40,13 +40,11 @@ public abstract class BaseYdbExecutor implements YdbExecutor {
     private final boolean useStreamResultSet;
 
     private final AtomicReference<YdbQueryResult> currResult;
-    protected final boolean traceEnabled;
     protected final String prefixPragma;
     protected final YdbTypes types;
 
     public BaseYdbExecutor(YdbContext ctx) {
         this.retryCtx = ctx.getRetryCtx();
-        this.traceEnabled = ctx.isTxTracerEnabled();
         this.sessionTimeout = ctx.getOperationProperties().getSessionTimeout();
         this.useStreamResultSet = ctx.getOperationProperties().getUseStreamResultSets();
         this.tableClient = ctx.getTableClient();

jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java

@@ -27,6 +27,7 @@
 import tech.ydb.jdbc.YdbTracer;
 import tech.ydb.jdbc.common.YdbTypes;
 import tech.ydb.jdbc.exception.ExceptionFactory;
+import tech.ydb.jdbc.impl.YdbTracerImpl;
 import tech.ydb.jdbc.impl.YdbTracerNone;
 import tech.ydb.jdbc.query.QueryType;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
@@ -90,6 +91,8 @@ public class YdbContext implements AutoCloseable {
     private final boolean autoResizeSessionPool;
     private final AtomicInteger connectionsCount = new AtomicInteger();
 
+    private YdbTracer.Storage tracerStorage;
+
     private YdbContext(
             YdbConfig config,
             YdbOperationProperties operationProperties,
@@ -137,12 +140,18 @@ private YdbContext(
             prefixPath = transport.getDatabase();
             prefixPragma = "";
         }
+
+        this.tracerStorage = config.isTxTracedEnabled() ? YdbTracerImpl.ENABLED : YdbTracerNone.DISABLED;
     }
 
     public YdbTypes getTypes() {
         return types;
     }
 
+    public void setTracerStorage(YdbTracer.Storage storage) {
+        this.tracerStorage = storage;
+    }
+
     /**
      * Grpc Transport for other API YDB server clients
      *
@@ -152,6 +161,10 @@ public GrpcTransport getGrpcTransport() {
         return grpcTransport;
     }
 
+    public YdbTracer getTracer() {
+        return tracerStorage.get();
+    }
+
     private String joined(String path1, String path2) {
         return path1.endsWith("/") || path2.startsWith("/") ? path1 + path2 : path1 + "/" + path2;
     }
@@ -188,14 +201,6 @@ public String getUsername() {
         return config.getUsername();
     }
 
-    public YdbTracer getTracer() {
-        return config.isTxTracedEnabled() ? YdbTracer.current() : YdbTracerNone.current();
-    }
-
-    public boolean isTxTracerEnabled() {
-        return config.isTxTracedEnabled();
-    }
-
     public YdbExecutor createExecutor() throws SQLException {
         if (config.isUseQueryService()) {
             return new QueryServiceExecutor(this, operationProps.getTransactionLevel(), operationProps.isAutoCommit());
@@ -264,6 +269,8 @@ public void register() {
     }
 
     public void deregister() {
+        tracerStorage.clear();
+
         int actual = connectionsCount.decrementAndGet();
         int maxSize = tableClient.sessionPoolStats().getMaxSize();
         if (autoResizeSessionPool && maxSize > SESSION_POOL_RESIZE_STEP) {
@@ -409,9 +416,9 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
             String tablePath = joined(getPrefixPath(), query.getYqlBatcher().getTableName());
             TableDescription description = tableDescribeCache.getIfPresent(tablePath);
             if (description == null) {
-                YdbTracer tracer = getTracer();
+                YdbTracer tracer = tracerStorage.get();
                 tracer.trace("--> describe table");
-                tracer.query(tablePath);
+                tracer.trace(tablePath);
 
                 DescribeTableSettings settings = withDefaultTimeout(new DescribeTableSettings());
                 Result<TableDescription> result = retryCtx.supplyResult(
@@ -449,9 +456,9 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
         Map<String, Type> queryTypes = queryParamsCache.getIfPresent(query.getOriginQuery());
         if (queryTypes == null) {
             String yql = prefixPragma + query.getPreparedYql();
-            YdbTracer tracer = getTracer();
+            YdbTracer tracer = tracerStorage.get();
             tracer.trace("--> prepare data query");
-            tracer.query(yql);
+            tracer.trace(yql);
 
             PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings());
             Result<DataQuery> result = retryCtx.supplyResult(

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java

@@ -26,7 +26,6 @@
 import tech.ydb.jdbc.YdbPrepareMode;
 import tech.ydb.jdbc.YdbPreparedStatement;
 import tech.ydb.jdbc.YdbStatement;
-import tech.ydb.jdbc.YdbTracer;
 import tech.ydb.jdbc.context.YdbContext;
 import tech.ydb.jdbc.context.YdbExecutor;
 import tech.ydb.jdbc.context.YdbValidator;
@@ -108,7 +107,6 @@ public void close() throws SQLException {
         validator.clearWarnings();
         executor.close();
         ctx.deregister();
-        YdbTracer.clear();
     }
 
     @Override

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerImpl.java

@@ -16,17 +16,34 @@
  */
 public class YdbTracerImpl implements YdbTracer {
     private static final Logger LOGGER = Logger.getLogger(YdbTracer.class.getName());
-    private static final ThreadLocal<YdbTracerImpl> LOCAL = new ThreadLocal<>();
+    private static final ThreadLocal<YdbTracer> LOCAL = new ThreadLocal<>();
     private static final AtomicLong ANONYMOUS_COUNTER = new AtomicLong(0);
 
+    public static final Storage ENABLED = new Storage() {
+        @Override
+        public YdbTracer get() {
+            YdbTracer tracer = LOCAL.get();
+            if (tracer == null) {
+                tracer = new YdbTracerImpl();
+                LOCAL.set(tracer);
+            }
+
+            return tracer;
+        }
+
+        @Override
+        public void clear() {
+            LOCAL.remove();
+        }
+    };
+
     private final Date startDate = new Date();
     private final long startedAt = System.currentTimeMillis();
     private final List<Record> records = new ArrayList<>();
 
     private String txID = null;
     private String label = null;
     private boolean isMarked = false;
-    private boolean isClosed = false;
 
     private class Record {
         private final long executedAt = System.currentTimeMillis();
@@ -39,20 +56,6 @@ private class Record {
         }
     }
 
-    public static void clear() {
-        LOCAL.remove();
-    }
-
-    public static YdbTracer current() {
-        YdbTracerImpl tracer = LOCAL.get();
-        if (tracer == null || tracer.isClosed) {
-            tracer = new YdbTracerImpl();
-            LOCAL.set(tracer);
-        }
-
-        return tracer;
-    }
-
     @Override
     public void trace(String message) {
         records.add(new Record(message, false));
@@ -82,8 +85,6 @@ public void markToPrint(String label) {
 
     @Override
     public void close() {
-        isClosed = true;
-
         LOCAL.remove();
 
         final Level level = isMarked ? Level.INFO : Level.FINE;

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerNone.java

@@ -7,11 +7,8 @@
  * @author Aleksandr Gorshenin
  */
 public class YdbTracerNone implements YdbTracer {
-    private static final YdbTracerNone INSTANCE = new YdbTracerNone();
-
-    public static YdbTracer current() {
-        return INSTANCE;
-    }
+    private static final YdbTracer NONE = new YdbTracerNone();
+    public static final YdbTracer.Storage DISABLED = () -> NONE;
 
     @Override
     public void trace(String message) { }