fix: Remove singleton access to request profiler. Use instance of hoo…

…k instead. (googleapis#2567) * Add profiler for request execution details. The usage of the new API will be added in the next PR * Add profiler for request execution details. The usage of the new API will be added in the next PR * add new code change * . * Remove directly access to singleton
yifatgortler · Jul 25, 2024 · 167f9ec · 167f9ec
1 parent 8f66b35
commit 167f9ec
Show file tree

Hide file tree

Showing 9 changed files with 195 additions and 158 deletions.
diff --git a/...-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java b/...-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java
@@ -247,6 +247,8 @@ class ConnectionWorker implements AutoCloseable {
    */
   private final RetrySettings retrySettings;
 
+  private final RequestProfiler.RequestProfilerHook requestProfilerHook;
+
   private static String projectMatching = "projects/[^/]+/";
   private static Pattern streamPatternProject = Pattern.compile(projectMatching);
 
@@ -386,7 +388,8 @@ public ConnectionWorker(
       String traceId,
       @Nullable String compressorName,
       BigQueryWriteSettings clientSettings,
-      RetrySettings retrySettings)
+      RetrySettings retrySettings,
+      boolean enableRequestProfiler)
       throws IOException {
     this.lock = new ReentrantLock();
     this.hasMessageInWaitingQueue = lock.newCondition();
@@ -410,6 +413,7 @@ public ConnectionWorker(
     this.compressorName = compressorName;
     this.retrySettings = retrySettings;
     this.telemetryAttributes = buildOpenTelemetryAttributes();
+    this.requestProfilerHook = new RequestProfiler.RequestProfilerHook(enableRequestProfiler);
     registerOpenTelemetryMetrics();
 
     // Always recreate a client for connection worker.
@@ -503,7 +507,7 @@ private boolean shouldWaitForBackoff(AppendRequestAndResponse requestWrapper) {
 
   private void waitForBackoffIfNecessary(AppendRequestAndResponse requestWrapper) {
     lock.lock();
-    RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+    requestProfilerHook.startOperation(
         RequestProfiler.OperationName.RETRY_BACKOFF, requestWrapper.requestUniqueId);
     try {
       Condition condition = lock.newCondition();
@@ -513,7 +517,7 @@ private void waitForBackoffIfNecessary(AppendRequestAndResponse requestWrapper)
     } catch (InterruptedException e) {
       throw new IllegalStateException(e);
     } finally {
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+      requestProfilerHook.endOperation(
           RequestProfiler.OperationName.RETRY_BACKOFF, requestWrapper.requestUniqueId);
       lock.unlock();
     }
@@ -535,7 +539,7 @@ private void addMessageToWaitingQueue(
     ++this.inflightRequests;
     this.inflightBytes += requestWrapper.messageSize;
     hasMessageInWaitingQueue.signal();
-    RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+    requestProfilerHook.startOperation(
         RequestProfiler.OperationName.WAIT_QUEUE, requestWrapper.requestUniqueId);
     if (addToFront) {
       waitingRequestQueue.addFirst(requestWrapper);
@@ -649,13 +653,12 @@ private ApiFuture<AppendRowsResponse> appendInternal(
                 writerId));
         return requestWrapper.appendResult;
       }
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
-          RequestProfiler.OperationName.WAIT_QUEUE, requestUniqueId);
+      requestProfilerHook.startOperation(RequestProfiler.OperationName.WAIT_QUEUE, requestUniqueId);
       ++this.inflightRequests;
       this.inflightBytes += requestWrapper.messageSize;
       waitingRequestQueue.addLast(requestWrapper);
       hasMessageInWaitingQueue.signal();
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+      requestProfilerHook.startOperation(
           RequestProfiler.OperationName.WAIT_INFLIGHT_QUOTA, requestUniqueId);
       try {
         maybeWaitForInflightQuota();
@@ -665,7 +668,7 @@ private ApiFuture<AppendRowsResponse> appendInternal(
         this.inflightBytes -= requestWrapper.messageSize;
         throw ex;
       }
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+      requestProfilerHook.endOperation(
           RequestProfiler.OperationName.WAIT_INFLIGHT_QUOTA, requestUniqueId);
       return requestWrapper.appendResult;
     } finally {
@@ -831,10 +834,10 @@ private void appendLoop() {
           while (!inflightRequestQueue.isEmpty()) {
             AppendRequestAndResponse requestWrapper = inflightRequestQueue.pollLast();
             // Consider the backend latency as completed for the current request.
-            RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+            requestProfilerHook.endOperation(
                 RequestProfiler.OperationName.RESPONSE_LATENCY, requestWrapper.requestUniqueId);
             requestWrapper.requestSendTimeStamp = null;
-            RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+            requestProfilerHook.startOperation(
                 RequestProfiler.OperationName.WAIT_QUEUE, requestWrapper.requestUniqueId);
             waitingRequestQueue.addFirst(requestWrapper);
           }
@@ -845,7 +848,7 @@ private void appendLoop() {
         }
         while (!this.waitingRequestQueue.isEmpty()) {
           AppendRequestAndResponse requestWrapper = this.waitingRequestQueue.pollFirst();
-          RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+          requestProfilerHook.endOperation(
               RequestProfiler.OperationName.WAIT_QUEUE, requestWrapper.requestUniqueId);
           waitForBackoffIfNecessary(requestWrapper);
           this.inflightRequestQueue.add(requestWrapper);
@@ -931,7 +934,7 @@ private void appendLoop() {
         }
         firstRequestForTableOrSchemaSwitch = false;
 
-        RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+        requestProfilerHook.startOperation(
             RequestProfiler.OperationName.RESPONSE_LATENCY, requestUniqueId);
 
         // Send should only throw an exception if there is a problem with the request. The catch
@@ -1212,7 +1215,7 @@ private void requestCallback(AppendRowsResponse response) {
       }
       if (!this.inflightRequestQueue.isEmpty()) {
         requestWrapper = pollFirstInflightRequestQueue();
-        RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+        requestProfilerHook.endOperation(
             RequestProfiler.OperationName.RESPONSE_LATENCY, requestWrapper.requestUniqueId);
       } else if (inflightCleanuped) {
         // It is possible when requestCallback is called, the inflight queue is already drained
@@ -1277,7 +1280,7 @@ private void requestCallback(AppendRowsResponse response) {
               requestWrapper.appendResult.set(response);
             }
           } finally {
-            RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+            requestProfilerHook.endOperation(
                 RequestProfiler.OperationName.TOTAL_LATENCY, requestWrapper.requestUniqueId);
           }
         });

diff --git a/...querystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java b/...querystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java
@@ -205,6 +205,8 @@ public abstract static class Builder {
   /** Static setting for connection pool. */
   private static Settings settings = Settings.builder().build();
 
+  private final boolean enableRequestProfiler;
+
   ConnectionWorkerPool(
       long maxInflightRequests,
       long maxInflightBytes,
@@ -213,7 +215,8 @@ public abstract static class Builder {
       String traceId,
       @Nullable String comperssorName,
       BigQueryWriteSettings clientSettings,
-      RetrySettings retrySettings) {
+      RetrySettings retrySettings,
+      boolean enableRequestProfiler) {
     this.maxInflightRequests = maxInflightRequests;
     this.maxInflightBytes = maxInflightBytes;
     this.maxRetryDuration = maxRetryDuration;
@@ -223,6 +226,7 @@ public abstract static class Builder {
     this.clientSettings = clientSettings;
     this.currentMaxConnectionCount = settings.minConnectionsPerRegion();
     this.retrySettings = retrySettings;
+    this.enableRequestProfiler = enableRequestProfiler;
   }
 
   /**
@@ -404,7 +408,8 @@ private ConnectionWorker createConnectionWorker(
             traceId,
             compressorName,
             clientSettings,
-            retrySettings);
+            retrySettings,
+            enableRequestProfiler);
     connectionWorkerPool.add(connectionWorker);
     log.info(
         String.format(

diff --git a/...d-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/RequestProfiler.java b/...d-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/RequestProfiler.java
@@ -80,7 +80,7 @@ enum OperationName {
   private static final int MAX_CACHED_REQUEST = 100000;
 
   // Singleton for easier access.
-  static final RequestProfiler REQUEST_PROFILER_SINGLETON = new RequestProfiler();
+  private static final RequestProfiler REQUEST_PROFILER_SINGLETON = new RequestProfiler();
 
   // Tunable static variable indicate how many top longest latency requests we should consider.
   private static final int DEFAULT_TOP_K = 20;
@@ -172,33 +172,35 @@ void flushAndPrintReport() {
   // Periodically trigger the report generation.
   void startPeriodicalReportFlushing() {
     this.enableProfiiler = true;
-    this.flushThread =
-        new Thread(
-            new Runnable() {
-              @Override
-              public void run() {
-                try {
-                  while (true) {
-                    try {
-                      TimeUnit.MILLISECONDS.sleep(FLUSH_PERIOD.toMillis());
-                    } catch (InterruptedException e) {
-                      log.warning("Flush report thread is interrupted by " + e.toString());
-                      throw new RuntimeException(e);
+    if (this.flushThread == null || !this.flushThread.isAlive()) {
+      this.flushThread =
+          new Thread(
+              new Runnable() {
+                @Override
+                public void run() {
+                  try {
+                    while (true) {
+                      try {
+                        TimeUnit.MILLISECONDS.sleep(FLUSH_PERIOD.toMillis());
+                      } catch (InterruptedException e) {
+                        log.warning("Flush report thread is interrupted by " + e.toString());
+                        throw new RuntimeException(e);
+                      }
+                      flushAndPrintReport();
                     }
-                    flushAndPrintReport();
+                  } catch (Exception ex) {
+                    // Mute any exception thrown from profiler process as we don't want to
+                    // interrupt normal operations.
+                    log.warning(
+                        "Exception thrown request profiler ignored, this is suggesting faulty "
+                            + "implementation of "
+                            + "RequestProfiler, exception context: "
+                            + ex.toString());
                   }
-                } catch (Exception ex) {
-                  // Mute any exception thrown from profiler process as we don't want to
-                  // interrupt normal operations.
-                  log.warning(
-                      "Exception thrown request profiler ignored, this is suggesting faulty "
-                          + "implementation of "
-                          + "RequestProfiler, exception context: "
-                          + ex.toString());
                 }
-              }
-            });
-    this.flushThread.start();
+              });
+      this.flushThread.start();
+    }
   }
 
   String flushAndGenerateReportText() {
@@ -402,7 +404,48 @@ void internalDisableAndClearProfiler() {
     FLUSH_PERIOD = DEFAULT_FLUSH_PERIOD;
   }
 
-  public static void disableAndClearProfiler() {
+  public static void disableAndResetProfiler() {
     REQUEST_PROFILER_SINGLETON.internalDisableAndClearProfiler();
   }
+
+  /**
+   * A hook for easier access to request profiler. Otherwise we have to trigger tedious if clauses
+   * to check whether profiler is enabled before every caller's trigger of the request profiler.
+   * This is because profiler is shared statically across instances.
+   */
+  static class RequestProfilerHook {
+    private boolean enableRequestProfiler = false;
+
+    RequestProfilerHook(boolean enableRequestProfiler) {
+      this.enableRequestProfiler = enableRequestProfiler;
+    }
+
+    // Mimic the api exposed by the main request profiler.
+    void startOperation(OperationName operationName, String requestUniqueId) {
+      if (this.enableRequestProfiler) {
+        RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(operationName, requestUniqueId);
+      }
+    }
+
+    // Mimic the api exposed by the main request profiler.
+    void endOperation(OperationName operationName, String requestUniqueId) {
+      if (this.enableRequestProfiler) {
+        RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(operationName, requestUniqueId);
+      }
+    }
+
+    void startPeriodicalReportFlushing() {
+      if (this.enableRequestProfiler) {
+        RequestProfiler.REQUEST_PROFILER_SINGLETON.startPeriodicalReportFlushing();
+      }
+    }
+
+    String flushAndGenerateReportText() {
+      return RequestProfiler.REQUEST_PROFILER_SINGLETON.flushAndGenerateReportText();
+    }
+
+    void enableProfiler() {
+      REQUEST_PROFILER_SINGLETON.enableProfiler();
+    }
+  }
 }
diff --git a/...rystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java b/...rystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java
@@ -71,6 +71,9 @@ public class SchemaAwareStreamWriter<T> implements AutoCloseable {
   // the user provides the table schema, we should always use that schema.
   private final boolean skipRefreshStreamWriter;
 
+  // Provide access to the request profiler.
+  private final RequestProfiler.RequestProfilerHook requestProfilerHook;
+
   /**
    * Constructs the SchemaAwareStreamWriter
    *
@@ -104,8 +107,10 @@ private SchemaAwareStreamWriter(Builder<T> builder)
         builder.defaultMissingValueInterpretation);
     streamWriterBuilder.setMissingValueInterpretationMap(builder.missingValueInterpretationMap);
     streamWriterBuilder.setClientId(builder.clientId);
+    streamWriterBuilder.setEnableLatencyProfiler(builder.enableRequestProfiler);
+    requestProfilerHook = new RequestProfiler.RequestProfilerHook(builder.enableRequestProfiler);
     if (builder.enableRequestProfiler) {
-      streamWriterBuilder.setEnableLatencyProfiler(builder.enableRequestProfiler);
+      requestProfilerHook.startPeriodicalReportFlushing();
     }
     this.streamWriter = streamWriterBuilder.build();
     this.streamName = builder.streamName;
@@ -128,12 +133,12 @@ private SchemaAwareStreamWriter(Builder<T> builder)
   public ApiFuture<AppendRowsResponse> append(Iterable<T> items)
       throws IOException, DescriptorValidationException {
     String requestUniqueId = generateRequestUniqueId();
-    RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+    requestProfilerHook.startOperation(
         RequestProfiler.OperationName.TOTAL_LATENCY, requestUniqueId);
     try {
       return appendWithUniqueId(items, -1, requestUniqueId);
     } catch (Exception ex) {
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+      requestProfilerHook.endOperation(
           RequestProfiler.OperationName.TOTAL_LATENCY, requestUniqueId);
       throw ex;
     }
@@ -198,12 +203,12 @@ private List<DynamicMessage> buildMessage(Iterable<T> items)
   public ApiFuture<AppendRowsResponse> append(Iterable<T> items, long offset)
       throws IOException, DescriptorValidationException {
     String requestUniqueId = generateRequestUniqueId();
-    RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+    requestProfilerHook.startOperation(
         RequestProfiler.OperationName.TOTAL_LATENCY, requestUniqueId);
     try {
       return appendWithUniqueId(items, offset, requestUniqueId);
     } catch (Exception ex) {
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+      requestProfilerHook.endOperation(
           RequestProfiler.OperationName.TOTAL_LATENCY, requestUniqueId);
       throw ex;
     }
@@ -214,7 +219,7 @@ ApiFuture<AppendRowsResponse> appendWithUniqueId(
       throws DescriptorValidationException, IOException {
     // Handle schema updates in a Thread-safe way by locking down the operation
     synchronized (this) {
-      RequestProfiler.REQUEST_PROFILER_SINGLETON.startOperation(
+      requestProfilerHook.startOperation(
           RequestProfiler.OperationName.JSON_TO_PROTO_CONVERSION, requestUniqueId);
       ProtoRows.Builder rowsBuilder = ProtoRows.newBuilder();
       try {
@@ -247,7 +252,7 @@ ApiFuture<AppendRowsResponse> appendWithUniqueId(
               rowIndexToErrorMessage);
         }
       } finally {
-        RequestProfiler.REQUEST_PROFILER_SINGLETON.endOperation(
+        requestProfilerHook.endOperation(
             RequestProfiler.OperationName.JSON_TO_PROTO_CONVERSION, requestUniqueId);
       }
       return this.streamWriter.appendWithUniqueId(rowsBuilder.build(), offset, requestUniqueId);
@@ -514,9 +519,6 @@ private Builder(
         this.skipRefreshStreamWriter = true;
       }
       this.toProtoConverter = toProtoConverter;
-      if (this.enableRequestProfiler) {
-        RequestProfiler.REQUEST_PROFILER_SINGLETON.startPeriodicalReportFlushing();
-      }
     }
 
     /**