Update workloadmeta protobuf

DataDog · Feb 7, 2025 · 22cd8f2 · 22cd8f2
1 parent fa33a49
commit 22cd8f2
Show file tree

Hide file tree

Showing 4 changed files with 137 additions and 38 deletions.
diff --git a/comp/core/workloadmeta/proto/proto.go b/comp/core/workloadmeta/proto/proto.go
@@ -134,19 +134,22 @@ func protoContainerFromWorkloadmetaContainer(container *workloadmeta.Container)
 		return nil, err
 	}
 
+	protoAllocatedResources := toProtoAllocatedResources(container.AllocatedResources)
+
 	return &pb.Container{
-		EntityId:      protoEntityID,
-		EntityMeta:    toProtoEntityMetaFromContainer(container),
-		EnvVars:       container.EnvVars,
-		Hostname:      container.Hostname,
-		Image:         toProtoImage(&container.Image),
-		NetworkIps:    container.NetworkIPs,
-		Pid:           int32(container.PID),
-		Ports:         pbContainerPorts,
-		Runtime:       protoRuntime,
-		State:         protoContainerState,
-		CollectorTags: container.CollectorTags,
-		CgroupPath:    container.CgroupPath,
+		EntityId:           protoEntityID,
+		EntityMeta:         toProtoEntityMetaFromContainer(container),
+		EnvVars:            container.EnvVars,
+		Hostname:           container.Hostname,
+		Image:              toProtoImage(&container.Image),
+		NetworkIps:         container.NetworkIPs,
+		Pid:                int32(container.PID),
+		Ports:              pbContainerPorts,
+		Runtime:            protoRuntime,
+		State:              protoContainerState,
+		CollectorTags:      container.CollectorTags,
+		CgroupPath:         container.CgroupPath,
+		AllocatedResources: protoAllocatedResources,
 	}, nil
 }
 
@@ -281,6 +284,18 @@ func toProtoContainerState(state *workloadmeta.ContainerState) (*pb.ContainerSta
 	return res, nil
 }
 
+func toProtoAllocatedResources(resources []workloadmeta.ContainerAllocatedResource) []*pb.ContainerAllocatedResource {
+	var protoAllocatedResources []*pb.ContainerAllocatedResource
+	for _, resource := range resources {
+		protoAllocatedResources = append(protoAllocatedResources, &pb.ContainerAllocatedResource{
+			Name: resource.Name,
+			ID:   resource.ID,
+		})
+	}
+
+	return protoAllocatedResources
+}
+
 func toProtoContainerStatus(status workloadmeta.ContainerStatus) (pb.ContainerStatus, error) {
 	switch status {
 	case "", workloadmeta.ContainerStatusUnknown:
@@ -600,19 +615,22 @@ func toWorkloadmetaContainer(protoContainer *pb.Container) (*workloadmeta.Contai
 		return nil, err
 	}
 
+	resources := toWorkloadmetaAllocatedResources(protoContainer.AllocatedResources)
+
 	return &workloadmeta.Container{
-		EntityID:      entityID,
-		EntityMeta:    toWorkloadmetaEntityMeta(protoContainer.EntityMeta),
-		EnvVars:       protoContainer.EnvVars,
-		Hostname:      protoContainer.Hostname,
-		Image:         toWorkloadmetaImage(protoContainer.Image),
-		NetworkIPs:    protoContainer.NetworkIps,
-		PID:           int(protoContainer.Pid),
-		Ports:         ports,
-		Runtime:       runtime,
-		State:         state,
-		CollectorTags: protoContainer.CollectorTags,
-		CgroupPath:    protoContainer.CgroupPath,
+		EntityID:           entityID,
+		EntityMeta:         toWorkloadmetaEntityMeta(protoContainer.EntityMeta),
+		EnvVars:            protoContainer.EnvVars,
+		Hostname:           protoContainer.Hostname,
+		Image:              toWorkloadmetaImage(protoContainer.Image),
+		NetworkIPs:         protoContainer.NetworkIps,
+		PID:                int(protoContainer.Pid),
+		Ports:              ports,
+		Runtime:            runtime,
+		State:              state,
+		CollectorTags:      protoContainer.CollectorTags,
+		CgroupPath:         protoContainer.CgroupPath,
+		AllocatedResources: resources,
 	}, nil
 }
 
@@ -624,6 +642,18 @@ func toWorkloadmetaContainerPort(protoPort *pb.ContainerPort) workloadmeta.Conta
 	}
 }
 
+func toWorkloadmetaAllocatedResources(protoAllocatedResources []*pb.ContainerAllocatedResource) []workloadmeta.ContainerAllocatedResource {
+	var resources []workloadmeta.ContainerAllocatedResource
+	for _, protoResource := range protoAllocatedResources {
+		resources = append(resources, workloadmeta.ContainerAllocatedResource{
+			Name: protoResource.Name,
+			ID:   protoResource.ID,
+		})
+	}
+
+	return resources
+}
+
 func toWorkloadmetaEntityID(protoEntityID *pb.WorkloadmetaEntityId) (workloadmeta.EntityID, error) {
 	kind, err := toWorkloadmetaKind(protoEntityID.Kind)
 	if err != nil {

diff --git a/comp/core/workloadmeta/proto/proto_test.go b/comp/core/workloadmeta/proto/proto_test.go
@@ -84,6 +84,9 @@ func TestConversions(t *testing.T) {
 					CollectorTags: []string{
 						"tag1",
 					},
+					AllocatedResources: []workloadmeta.ContainerAllocatedResource{
+						{Name: "nvidia.com/gpu", ID: "gpu1"},
+					},
 				},
 			},
 			protoWorkloadmetaEvent: &pb.WorkloadmetaEvent{
@@ -138,6 +141,9 @@ func TestConversions(t *testing.T) {
 					CollectorTags: []string{
 						"tag1",
 					},
+					AllocatedResources: []*pb.ContainerAllocatedResource{
+						{Name: "nvidia.com/gpu", ID: "gpu1"},
+					},
 				},
 			},
 			expectsError: false,

diff --git a/pkg/proto/datadog/workloadmeta/workloadmeta.proto b/pkg/proto/datadog/workloadmeta/workloadmeta.proto
@@ -95,6 +95,11 @@ message ContainerState {
   int64 exitCode = 7;
 }
 
+message ContainerAllocatedResource {
+  string Name = 1;
+  string ID = 2;
+}
+
 message Container {
   WorkloadmetaEntityId entityId = 1;
   EntityMeta entityMeta = 2;
@@ -108,6 +113,7 @@ message Container {
   ContainerState state = 10;
   repeated string collectorTags = 11;
   string cgroupPath = 12;
+  repeated ContainerAllocatedResource allocatedResources = 13;
 }
 
 message KubernetesPodOwner {

diff --git a/pkg/proto/pbgo/core/workloadmeta.pb.go b/pkg/proto/pbgo/core/workloadmeta.pb.go