Log bytes written to disk when client app shuts down

Author: cheb0

ulyp-agent-core/src/main/java/com/ulyp/agent/AgentShutdownHook.java

@@ -1,6 +1,7 @@
 package com.ulyp.agent;
 
 import com.ulyp.agent.queue.RecordingEventQueue;
+import com.ulyp.core.util.ByteSize;
 import com.ulyp.storage.writer.RecordingDataWriter;
 import lombok.extern.slf4j.Slf4j;
 
@@ -17,6 +18,8 @@ public void run() {
             return;
         }
 
+        log.info("Shutting down the agent conext. Will wait for recording data to be flushed to disk...");
+
         RecordingEventQueue recordingEventQueue = ctx.getRecordingEventQueue();
         try {
             recordingEventQueue.sync(Duration.ofSeconds(30));
@@ -35,6 +38,8 @@ public void run() {
         } catch (TimeoutException e) {
             // ignore
         }
+        long totalBytesWritten = storageWriter.estimateBytesWritten();
+        log.info("Total bytes written to file: {}", ByteSize.toHumanReadable(totalBytesWritten));
         storageWriter.close();
     }
 }

ulyp-storage/src/main/java/com/ulyp/storage/writer/PerTypeStats.java

@@ -18,14 +18,6 @@ public void reset() {
 
     }
 
-    public long getTotalCount() {
-        return totalCount.get();
-    }
-
-    public void addToCount(long delta) {
-        totalCount.addAndGet(delta);
-    }
-
     public long getTotalBytes() {
         return totalBytes.getByteSize();
     }

ulyp-storage/src/main/java/com/ulyp/storage/writer/StatsRecordingDataWriter.java

@@ -12,22 +12,22 @@
 
 import java.time.Duration;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
 
 @Slf4j
 public class StatsRecordingDataWriter implements RecordingDataWriter {
 
     private final RecordingDataWriter delegate;
-    private final BytesCounter typeBytes;
-    private final BytesCounter methodBytes;
-    private final BytesCounter callsBytes;
-
-    private final PerTypeStats totalBytesWritten = new PerTypeStats("Total bytes written");
+    private final BytesCounter typeBytesWritten;
+    private final BytesCounter methodBytesWritten;
+    private final BytesCounter methodCallBytesWritten;
+    private final AtomicLong totalBytesWritten = new AtomicLong(0L);
 
     public StatsRecordingDataWriter(Metrics metrics, RecordingDataWriter delegate) {
         this.delegate = delegate;
-        this.typeBytes = metrics.getOrCreateByteCounter("writer.bytes.types");
-        this.methodBytes = metrics.getOrCreateByteCounter("writer.bytes.methods");
-        this.callsBytes = metrics.getOrCreateByteCounter("writer.bytes.calls");
+        this.typeBytesWritten = metrics.getOrCreateByteCounter("writer.bytes.types");
+        this.methodBytesWritten = metrics.getOrCreateByteCounter("writer.bytes.methods");
+        this.methodCallBytesWritten = metrics.getOrCreateByteCounter("writer.bytes.calls");
     }
 
     @Override
@@ -52,28 +52,28 @@ public void write(RecordingMetadata recordingMetadata) throws StorageException {
 
     @Override
     public void write(SerializedTypeList types) throws StorageException {
-        totalBytesWritten.addBytes(types.byteLength());
+        totalBytesWritten.addAndGet(types.byteLength());
         delegate.write(types);
-        typeBytes.add(types.byteLength(), types.size());
+        typeBytesWritten.add(types.byteLength(), types.size());
     }
 
     @Override
     public void write(SerializedRecordedMethodCallList callRecords) throws StorageException {
-        totalBytesWritten.addBytes(callRecords.bytesWritten());
+        totalBytesWritten.addAndGet(callRecords.bytesWritten());
         delegate.write(callRecords);
-        callsBytes.add(callRecords.bytesWritten(), callRecords.size());
+        methodCallBytesWritten.add(callRecords.bytesWritten(), callRecords.size());
     }
 
     @Override
     public long estimateBytesWritten() {
-        return totalBytesWritten.getTotalBytes();
+        return totalBytesWritten.get();
     }
 
     @Override
     public void write(SerializedMethodList methods) throws StorageException {
-        totalBytesWritten.addBytes(methods.byteLength());
+        totalBytesWritten.addAndGet(methods.byteLength());
         delegate.write(methods);
-        methodBytes.add(methods.byteLength(), methods.size());
+        methodBytesWritten.add(methods.byteLength(), methods.size());
     }
 
     @Override