Periodically capture resource metrics for every sandbox

packages/api/internal/handlers/sanbox_metrics.go

@@ -0,0 +1,101 @@
+package handlers
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"slices"
+	"strings"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/grafana/loki/pkg/loghttp"
+	"github.com/grafana/loki/pkg/logproto"
+	"go.opentelemetry.io/otel/attribute"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
+	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
+	"github.com/e2b-dev/infra/packages/api/internal/utils"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+func (a *APIStore) GetSandboxesSandboxIDMetrics(
+	c *gin.Context,
+	sandboxID string,
+) {
+	ctx := c.Request.Context()
+	sandboxID = utils.ShortID(sandboxID)
+
+	teamID := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team.ID
+
+	telemetry.SetAttributes(ctx,
+		attribute.String("instance.id", sandboxID),
+		attribute.String("team.id", teamID.String()),
+	)
+
+	var start time.Time
+
+	end := time.Now()
+
+	// Sanitize ID
+	// https://grafana.com/blog/2021/01/05/how-to-escape-special-characters-with-lokis-logql/
+	id := strings.ReplaceAll(sandboxID, "`", "")
+	query := fmt.Sprintf(
+		"{source=\"logs-collector\", category=\"metrics\", service=\"envd\", teamID=`%s`, sandboxID=`%s`}", teamID.String(), id)
+
+	res, err := a.lokiClient.QueryRange(query, 100, start, end, logproto.FORWARD, time.Duration(0), time.Duration(0), true)
+	if err != nil {
+		errMsg := fmt.Errorf("error when returning metrics for sandbox: %w", err)
+		telemetry.ReportCriticalError(ctx, errMsg)
+		a.sendAPIStoreError(c, http.StatusNotFound, fmt.Sprintf("Error returning metrics for sandbox '%s'", sandboxID))
+
+		return
+	}
+
+	switch res.Data.Result.Type() {
+	case loghttp.ResultTypeStream:
+		value := res.Data.Result.(loghttp.Streams)
+
+		metrics := make([]api.SandboxMetric, 0)
+
+		for _, stream := range value {
+			for _, entry := range stream.Entries {
+
+				var metric struct {
+					CPUPct      *float32 `json:"cpuPct"`
+					MemTotalMiB int64    `json:"memTotalMiB"`
+					MemUsedMiB  int64    `json:"memUsedMiB"`
+				}
+
+				err := json.Unmarshal([]byte(entry.Line), &metric)
+				if err != nil {
+					telemetry.ReportCriticalError(ctx, fmt.Errorf("failed to unmarshal metric: %w", err))
+					continue
+				}
+				metrics = append(metrics, api.SandboxMetric{
+					Timestamp:   entry.Timestamp,
+					CpuPct:      metric.CPUPct,
+					MemMiBTotal: metric.MemTotalMiB,
+					MemMiBUsed:  metric.MemUsedMiB,
+				})
+			}
+		}
+
+		// Sort metrics by timestamp (they are returned by the time they arrived in Loki)
+		slices.SortFunc(metrics, func(a, b api.SandboxMetric) int {
+			return a.Timestamp.Compare(b.Timestamp)
+		})
+
+		c.JSON(http.StatusOK, &api.SandboxMetrics{
+			Metrics: metrics,
+		})
+
+	default:
+		errMsg := fmt.Errorf("unexpected value type %T", res.Data.Result.Type())
+		telemetry.ReportCriticalError(ctx, errMsg)
+		a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error returning metrics for sandbox '%s", sandboxID))
+
+		return
+	}
+}

packages/envd/internal/host/metrics.go

@@ -9,9 +9,10 @@ import (
 )
 
 type Metrics struct {
-	CPU       float64 `json:"cpu_pct"`   // Percent rounded to 2 decimal places
-	Mem       uint64  `json:"mem_bytes"` // Total virtual memory in bytes
-	Timestamp int64   `json:"ts"`        // Unix Timestamp in UTC
+	Timestamp   int64   `json:"ts"`            // Unix Timestamp in UTC
+	CPUPercent  float64 `json:"cpu_pct"`       // Percent rounded to 2 decimal places
+	MemTotalMiB uint64  `json:"mem_total_mib"` // Total virtual memory in MiB
+	MemUsedMiB  uint64  `json:"mem_used_mib"`  // Used virtual memory in MiB
 }
 
 func GetMetrics() (*Metrics, error) {
@@ -20,6 +21,9 @@ func GetMetrics() (*Metrics, error) {
 		return nil, err
 	}
 
+	memUsedMiB := v.Used / 1024 / 1024
+	memTotalMiB := v.Total / 1024 / 1024
+
 	cpuPcts, err := cpu.Percent(0, false)
 	if err != nil {
 		return nil, err
@@ -32,8 +36,9 @@ func GetMetrics() (*Metrics, error) {
 	}
 
 	return &Metrics{
-		CPU:       cpuPctRounded,
-		Mem:       v.Total,
-		Timestamp: time.Now().UTC().Unix(),
+		Timestamp:   time.Now().UTC().Unix(),
+		CPUPercent:  cpuPctRounded,
+		MemUsedMiB:  memUsedMiB,
+		MemTotalMiB: memTotalMiB,
 	}, nil
 }

packages/orchestrator/internal/sandbox/checks.go

@@ -2,37 +2,49 @@ package sandbox
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"time"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+	"golang.org/x/mod/semver"
+)
+
+const (
+	healthCheckInterval      = 10 * time.Second
+	metricsCheckInterval     = 5 * time.Second
+	minEnvdVersionForMetrcis = "0.1.3"
 )
 
 func (s *Sandbox) logHeathAndUsage(ctx *utils.LockableCancelableContext) {
+	healthTicker := time.NewTicker(healthCheckInterval)
+	metricsTicker := time.NewTicker(metricsCheckInterval)
+	defer func() {
+		healthTicker.Stop()
+		metricsTicker.Stop()
+	}()
+
+	// Get metrics on sandbox startup
+	go s.LogMetrics(ctx)
+
 	ticker := time.NewTicker(15 * time.Second)
 	defer ticker.Stop()
 
 	for {
 		select {
-		case <-ticker.C:
+		case <-healthTicker.C:
 			childCtx, cancel := context.WithTimeout(ctx, time.Second)
 
 			ctx.Lock()
 			s.Healthcheck(childCtx, false)
 			ctx.Unlock()
 
 			cancel()
-
-			stats, err := s.stats.Stats()
-			if err != nil {
-				s.Logger.Warnf("failed to get stats: %s", err)
-			} else {
-				s.Logger.CPUUsage(stats.CPUCount)
-				s.Logger.MemoryUsage(stats.MemoryMB)
-			}
+		case <-metricsTicker.C:
+			s.LogMetrics(ctx)
 		case <-ctx.Done():
 			return
 		}
@@ -68,3 +80,55 @@ func (s *Sandbox) Healthcheck(ctx context.Context, alwaysReport bool) {
 		return
 	}
 }
+
+func (s *Sandbox) GetMetrics(ctx context.Context) (SandboxMetrics, error) {
+	address := fmt.Sprintf("http://%s:%d/metrics", s.slot.HostIP(), consts.DefaultEnvdServerPort)
+
+	request, err := http.NewRequestWithContext(ctx, "GET", address, nil)
+	if err != nil {
+		return SandboxMetrics{}, err
+	}
+
+	response, err := httpClient.Do(request)
+	if err != nil {
+		return SandboxMetrics{}, err
+	}
+	defer response.Body.Close()
+
+	if response.StatusCode != http.StatusOK {
+		err = fmt.Errorf("unexpected status code: %d", response.StatusCode)
+		return SandboxMetrics{}, err
+	}
+
+	var metrics SandboxMetrics
+	err = json.NewDecoder(response.Body).Decode(&metrics)
+	if err != nil {
+		return SandboxMetrics{}, err
+	}
+
+	return metrics, nil
+}
+
+func (s *Sandbox) LogMetrics(ctx context.Context) {
+	if isGTEVersion(s.Sandbox.EnvdVersion, minEnvdVersionForMetrcis) {
+		metrics, err := s.GetMetrics(ctx)
+		if err != nil {
+			s.Logger.Warnf("failed to get metrics: %s", err)
+		} else {
+			s.Logger.CPUPct(metrics.CPUPercent)
+			s.Logger.MemMiB(metrics.MemTotalMiB, metrics.MemUsedMiB)
+		}
+	}
+}
+
+func isGTEVersion(curVersion, minVersion string) bool {
+	if len(curVersion) > 0 && curVersion[0] != 'v' {
+		curVersion = "v" + curVersion
+	}
+
+	if !semver.IsValid(curVersion) {
+		return false
+	}
+
+	return semver.Compare(curVersion, minVersion) >= 0
+}

packages/orchestrator/internal/sandbox/metrics.go

@@ -0,0 +1,8 @@
+package sandbox
+
+type SandboxMetrics struct {
+	Timestamp   int64   `json:"ts"`            // Unix Timestamp in UTC
+	CPUPercent  float64 `json:"cpu_pct"`       // Percent rounded to 2 decimal places
+	MemTotalMiB uint64  `json:"mem_total_mib"` // Total virtual memory in MiB
+	MemUsedMiB  uint64  `json:"mem_used_mib"`  // Used virtual memory in MiB
+}

packages/orchestrator/internal/server/sandboxes.go

@@ -172,7 +172,18 @@ func (s *server) Delete(ctx context.Context, in *orchestrator.SandboxDeleteReque
 	// 	Ideally we would rely only on the goroutine defer.
 	s.sandboxes.Remove(in.SandboxId)
 
+	// Don't allow connecting to the sandbox anymore.
+	s.dns.Remove(in.SandboxId, sbx.Slot.HostIP())
+
+	// Remove the sandbox from the cache to prevent loading it again in API during the time the instance is stopping.
+	// Old comment:
+	// 	Ensure the sandbox is removed from cache.
+	// 	Ideally we would rely only on the goroutine defer.
+	s.sandboxes.Remove(in.SandboxId)
+
+	// Check health metrics before stopping the sandbox
 	sbx.Healthcheck(ctx, true)
+	sbx.LogMetrics(ctx)
 
 	err := sbx.Stop()
 	if err != nil {