feat: add inlight time metric

vinted · Mar 6, 2024 · 8d89378 · 8d89378
1 parent 672c917
commit 8d89378
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 1 deletion.
diff --git a/src/main/java/com/vinted/flink/bigquery/client/BigQueryProtoClientProvider.java b/src/main/java/com/vinted/flink/bigquery/client/BigQueryProtoClientProvider.java
@@ -55,6 +55,7 @@ public BigQueryStreamWriter<A> getWriter(String streamName, TableId table, RowVa
                     .setExecutorProvider(executorProvider)
                     .setLocation(table.getProject())
                     .setWriterSchema(protoSchema);
+
             StreamWriter.setMaxRequestCallbackWaitTime(this.writerSettings.getMaxRequestWaitCallbackTime());
             return new ProtoStreamWriter<>(serializer, streamWriterBuilder.build());
         } catch (IOException | Descriptors.DescriptorValidationException e) {

diff --git a/src/main/java/com/vinted/flink/bigquery/client/BigQueryStreamWriter.java b/src/main/java/com/vinted/flink/bigquery/client/BigQueryStreamWriter.java
@@ -11,5 +11,7 @@ public interface BigQueryStreamWriter<T> extends AutoCloseable {
     String getStreamName();
 
     String getWriterId();
+
+    long getInflightWaitSeconds();
     boolean isClosed();
 }
diff --git a/src/main/java/com/vinted/flink/bigquery/client/JsonStreamWriter.java b/src/main/java/com/vinted/flink/bigquery/client/JsonStreamWriter.java
@@ -42,6 +42,11 @@ public ApiFuture<AppendRowsResponse> append(Rows<A> data, long offset) {
         }
     }
 
+    @Override
+    public long getInflightWaitSeconds() {
+        return writer.getInflightWaitSeconds();
+    }
+
     @Override
     public String getStreamName() {
         return writer.getStreamName();

diff --git a/src/main/java/com/vinted/flink/bigquery/client/ProtoStreamWriter.java b/src/main/java/com/vinted/flink/bigquery/client/ProtoStreamWriter.java
@@ -38,6 +38,11 @@ public ApiFuture<AppendRowsResponse> append(Rows<A> data, long offset) {
         return writer.append(prows, offset);
     }
 
+    @Override
+    public long getInflightWaitSeconds() {
+        return writer.getInflightWaitSeconds();
+    }
+
     @Override
     public String getStreamName() {
         return writer.getStreamName();

diff --git a/src/main/java/com/vinted/flink/bigquery/sink/BigQuerySinkWriter.java b/src/main/java/com/vinted/flink/bigquery/sink/BigQuerySinkWriter.java
@@ -69,12 +69,22 @@ public BigQuerySinkWriter(
 
     }
 
+    private void registerInflightMetric(String streamName, BigQueryStreamWriter<A> writer) {
+        var group = metricGroup
+                .addGroup("stream", streamName)
+                .addGroup("writer_id", writer.getWriterId());
+
+        group.gauge("inflight_wait_seconds", writer::getInflightWaitSeconds);
+    }
+
     protected final BigQueryStreamWriter<A> streamWriter(String traceId, String streamName, TableId table) {
         var streamWithIndex = String.format("%s-%s",streamName, streamIndexIterator.next());
         return streamMap.computeIfAbsent(streamWithIndex, name -> {
             logger.trace("Trace-id {} Stream not found {}. Creating new stream", traceId, streamWithIndex);
             // Stream name can't contain index
-            return this.clientProvider.getWriter(streamName, table, rowSerializer);
+            var writer = this.clientProvider.getWriter(streamName, table, rowSerializer);
+            registerInflightMetric(streamName, writer);
+            return writer;
         });
     }
 
@@ -91,6 +101,7 @@ protected final void recreateStreamWriter(String traceId, String streamName, Str
                         logger.trace("Trace-id {} Could not close writer for {}", traceId, streamName);
                     }
                     newWriter = this.clientProvider.getWriter(streamName, table, rowSerializer);
+                    registerInflightMetric(streamName, writer);
                 }
                 return newWriter;
             });