Add disk usage stats to ClickHouse (#7905)

This is needed so that if we add a default task size for disk usage, we can analyze historical executions to see what the default should be.
buildbuddy-io · Nov 15, 2024 · a665327 · a665327
1 parent 513710c
commit a665327
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 7 deletions.
diff --git a/enterprise/server/remote_execution/execution_server/execution_server.go b/enterprise/server/remote_execution/execution_server/execution_server.go
@@ -329,14 +329,14 @@ func (s *ExecutionServer) updateExecution(ctx context.Context, executionID strin
 	}
 
 	if stage == repb.ExecutionStage_COMPLETED {
-		if err := s.recordExecution(ctx, executionID); err != nil {
+		if err := s.recordExecution(ctx, executionID, executeResponse.GetResult().GetExecutionMetadata()); err != nil {
 			log.CtxErrorf(ctx, "failed to record execution %q: %s", executionID, err)
 		}
 	}
 	return dbErr
 }
 
-func (s *ExecutionServer) recordExecution(ctx context.Context, executionID string) error {
+func (s *ExecutionServer) recordExecution(ctx context.Context, executionID string, md *repb.ExecutedActionMetadata) error {
 	if s.env.GetExecutionCollector() == nil || !olapdbconfig.WriteExecutionsToOLAPDBEnabled() {
 		return nil
 	}
@@ -361,9 +361,12 @@ func (s *ExecutionServer) recordExecution(ctx context.Context, executionID strin
 	rmd := bazel_request.GetRequestMetadata(ctx)
 	for _, link := range links {
 		executionProto := execution.TableExecToProto(&executionPrimaryDB, link)
-		if rmd != nil {
-			executionProto.TargetLabel = rmd.GetTargetId()
-		}
+		// Set fields that aren't stored in the primary DB
+		executionProto.TargetLabel = rmd.GetTargetId()
+		executionProto.DiskBytesRead = md.GetUsageStats().GetCgroupIoStats().GetRbytes()
+		executionProto.DiskBytesWritten = md.GetUsageStats().GetCgroupIoStats().GetWbytes()
+		executionProto.DiskWriteOperations = md.GetUsageStats().GetCgroupIoStats().GetWios()
+		executionProto.DiskReadOperations = md.GetUsageStats().GetCgroupIoStats().GetRios()
 		inv, err := s.env.GetExecutionCollector().GetInvocation(ctx, link.GetInvocationId())
 		if err != nil {
 			log.CtxErrorf(ctx, "failed to get invocation %q from ExecutionCollector: %s", link.GetInvocationId(), err)

diff --git a/proto/remote_execution.proto b/proto/remote_execution.proto
@@ -2407,6 +2407,10 @@ message StoredExecution {
   // UsageStats
   int64 peak_memory_bytes = 15;
   int64 cpu_nanos = 16;
+  int64 disk_bytes_read = 34;
+  int64 disk_bytes_written = 35;
+  int64 disk_read_operations = 36;
+  int64 disk_write_operations = 37;
 
   // Task Sizing
   int64 estimated_memory_bytes = 17;

diff --git a/server/util/clickhouse/clickhouse.go b/server/util/clickhouse/clickhouse.go
@@ -213,6 +213,10 @@ func buildExecution(in *repb.StoredExecution, inv *sipb.StoredInvocation) *schem
 		FileUploadDurationUsec:             in.GetFileUploadDurationUsec(),
 		PeakMemoryBytes:                    in.GetPeakMemoryBytes(),
 		CPUNanos:                           in.GetCpuNanos(),
+		DiskBytesRead:                      in.GetDiskBytesRead(),
+		DiskBytesWritten:                   in.GetDiskBytesWritten(),
+		DiskReadOperations:                 in.GetDiskReadOperations(),
+		DiskWriteOperations:                in.GetDiskWriteOperations(),
 		EstimatedMemoryBytes:               in.GetEstimatedMemoryBytes(),
 		EstimatedMilliCPU:                  in.GetEstimatedMilliCpu(),
 		QueuedTimestampUsec:                in.GetQueuedTimestampUsec(),

diff --git a/server/util/clickhouse/schema/schema.go b/server/util/clickhouse/schema/schema.go
@@ -171,8 +171,12 @@ type Execution struct {
 	FileUploadDurationUsec   int64
 
 	// UsageStats
-	PeakMemoryBytes int64
-	CPUNanos        int64
+	PeakMemoryBytes     int64
+	CPUNanos            int64
+	DiskBytesRead       int64
+	DiskBytesWritten    int64
+	DiskReadOperations  int64
+	DiskWriteOperations int64
 
 	// Task sizing
 	EstimatedMemoryBytes int64
@@ -248,6 +252,10 @@ func (e *Execution) AdditionalFields() []string {
 		"Tags",
 		"OutputPath",
 		"TargetLabel",
+		"DiskBytesRead",
+		"DiskBytesWritten",
+		"DiskReadOperations",
+		"DiskWriteOperations",
 	}
 }