replace use of instance client id in places where placement node id is needed

Author: sitole

packages/api/internal/cache/instance/sync.go

@@ -65,7 +65,7 @@ func (c *InstanceCache) Sync(ctx context.Context, instances []*InstanceInfo, nod
 
 	// Delete instances that are not in Orchestrator anymore
 	for _, item := range c.cache.Items() {
-		if item.Instance.ClientID != nodeID {
+		if item.Node.ID != nodeID {
 			continue
 		}
 		if time.Since(item.StartTime) <= syncSandboxRemoveGracePeriod {

packages/api/internal/handlers/sandbox_create.go

@@ -170,8 +170,6 @@ func (a *APIStore) PostSandboxes(c *gin.Context) {
 		return
 	}
 
-	c.Set("nodeID", sbx.ClientID)
-
 	c.JSON(http.StatusCreated, &sbx)
 }
 

packages/api/internal/orchestrator/cache.go

@@ -311,11 +311,10 @@ func (o *Orchestrator) getDeleteInstanceFunction(
 			duration,
 		)
 
-		node := o.GetNode(info.Instance.ClientID)
+		node := o.GetNode(info.Node.ID)
 		if node == nil {
-			zap.L().Error("failed to get node", zap.String("node_id", info.Instance.ClientID))
-
-			return fmt.Errorf("node '%s' not found", info.Instance.ClientID)
+			zap.L().Error("failed to get node", logger.WithNodeID(info.Node.ID))
+			return fmt.Errorf("node '%s' not found", info.Node.ID)
 		}
 
 		node.CPUUsage.Add(-info.VCpu)
@@ -324,9 +323,8 @@ func (o *Orchestrator) getDeleteInstanceFunction(
 		o.dns.Remove(ctx, info.Instance.SandboxID, node.Info.IPAddress)
 
 		if node.Client == nil {
-			zap.L().Error("client for node not found", zap.String("node_id", info.Instance.ClientID))
-
-			return fmt.Errorf("client for node '%s' not found", info.Instance.ClientID)
+			zap.L().Error("client for node not found", logger.WithNodeID(info.Node.ID))
+			return fmt.Errorf("client for node '%s' not found", info.Node.ID)
 		}
 
 		if ct == ClosePause {
@@ -416,9 +414,9 @@ func (o *Orchestrator) getInsertInstanceFunction(parentCtx context.Context, time
 			zap.Bool("auto_pause", info.AutoPause.Load()),
 		)
 
-		node := o.GetNode(info.Instance.ClientID)
+		node := o.GetNode(info.Node.ID)
 		if node == nil {
-			zap.L().Error("failed to get node", zap.String("node_id", info.Instance.ClientID))
+			zap.L().Error("failed to get node", logger.WithNodeID(info.Node.ID))
 		} else {
 			node.CPUUsage.Add(info.VCpu)
 			node.RamUsage.Add(info.RamMB)

packages/api/internal/orchestrator/keep_alive.go

@@ -17,7 +17,7 @@ func (o *Orchestrator) KeepAliveFor(ctx context.Context, sandboxID string, durat
 		return apiErr
 	}
 
-	err := o.UpdateSandbox(ctx, sbx.Instance.SandboxID, sbx.GetEndTime(), sbx.Instance.ClientID)
+	err := o.UpdateSandbox(ctx, sbx.Instance.SandboxID, sbx.GetEndTime(), sbx.Node.ID)
 	if err != nil {
 		zap.L().Error("Error when setting sandbox timeout", zap.Error(err), logger.WithSandboxID(sandboxID))
 		return &api.APIError{Code: http.StatusInternalServerError, ClientMsg: "Error when setting sandbox timeout", Err: err}

packages/api/internal/orchestrator/list_instances.go

@@ -71,7 +71,7 @@ func (o *Orchestrator) getSandboxes(ctx context.Context, node *nNode.NodeInfo) (
 					SandboxID:  config.SandboxId,
 					TemplateID: config.TemplateId,
 					Alias:      config.Alias,
-					ClientID:   node.ID, // to prevent mismatch use the node ID which we use for the request
+					ClientID:   sbx.ClientId,
 				},
 				config.ExecutionId,
 				&teamID,

packages/api/internal/orchestrator/node.go

@@ -20,6 +20,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	orchestratorinfo "github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator-info"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
@@ -183,9 +184,9 @@ func (o *Orchestrator) GetNodes() []*api.Node {
 	}
 
 	for _, sbx := range o.instanceCache.Items() {
-		n, ok := nodes[sbx.Instance.ClientID]
+		n, ok := nodes[sbx.Node.ID]
 		if !ok {
-			zap.L().Error("node for sandbox wasn't found", zap.String("client_id", sbx.Instance.ClientID), zap.String("sandbox_id", sbx.Instance.SandboxID))
+			zap.L().Error("node for sandbox wasn't found", logger.WithNodeID(sbx.Node.ID), logger.WithSandboxID(sbx.Instance.SandboxID))
 			continue
 		}
 
@@ -232,7 +233,7 @@ func (o *Orchestrator) GetNodeDetail(nodeID string) *api.NodeDetail {
 	}
 
 	for _, sbx := range o.instanceCache.Items() {
-		if sbx.Instance.ClientID == nodeID {
+		if sbx.Node.ID == nodeID {
 			var metadata *api.SandboxMetadata
 			if sbx.Metadata != nil {
 				meta := api.SandboxMetadata(sbx.Metadata)

packages/api/internal/orchestrator/pause_instance.go

@@ -85,9 +85,9 @@ func snapshotInstance(ctx context.Context, orch *Orchestrator, sbx *instance.Ins
 	_, childSpan := orch.tracer.Start(ctx, "snapshot-instance")
 	defer childSpan.End()
 
-	client, clientMd, err := orch.GetClient(sbx.Instance.ClientID)
+	client, clientMd, err := orch.GetClient(sbx.Node.ID)
 	if err != nil {
-		return fmt.Errorf("failed to get client '%s': %w", sbx.Instance.ClientID, err)
+		return fmt.Errorf("failed to get client '%s': %w", sbx.Node.ID, err)
 	}
 
 	reqCtx := metadata.NewOutgoingContext(ctx, clientMd)