Add traces inside endpoint autoscaler

pkg/abstractions/common/instance.go

@@ -178,6 +178,7 @@ func (i *AutoscaledInstance) ConsumeContainerEvent(event types.ContainerEvent) {
 }
 
 func (i *AutoscaledInstance) Monitor() error {
+
 	go i.Autoscaler.Start(i.Ctx) // Start the autoscaler
 	ignoreScalingEventWindow := time.Now().Add(-IgnoreScalingEventInterval)
 
@@ -263,6 +264,10 @@ func (i *AutoscaledInstance) HandleScalingEvent(desiredContainers int) error {
 }
 
 func (i *AutoscaledInstance) State() (*AutoscaledInstanceState, error) {
+	trace := common.TraceFunc(i.Ctx, "pkg/abstractions/common", "AutoscaledInstance.State",
+		attribute.String("stub.id", i.Stub.ExternalId))
+	defer trace.End()
+
 	containers, err := i.ContainerRepo.GetActiveContainersByStubId(i.Stub.ExternalId)
 	if err != nil {
 		return nil, err

pkg/abstractions/endpoint/autoscaler.go

@@ -1,9 +1,12 @@
 package endpoint
 
 import (
+	"fmt"
 	"math"
 
 	abstractions "github.com/beam-cloud/beta9/pkg/abstractions/common"
+	"github.com/beam-cloud/beta9/pkg/common"
+	"go.opentelemetry.io/otel/attribute"
 )
 
 type endpointAutoscalerSample struct {
@@ -12,16 +15,22 @@ type endpointAutoscalerSample struct {
 }
 
 func endpointSampleFunc(i *endpointInstance) (*endpointAutoscalerSample, error) {
+	trace := common.TraceFunc(i.Ctx, "pkg/abstractions/endpoint", "endpointSampleFunc",
+		attribute.String("stub.id", i.Stub.ExternalId))
+	defer trace.End()
+
 	totalRequests, err := i.TaskRepo.TasksInFlight(i.Ctx, i.Workspace.Name, i.Stub.ExternalId)
 	if err != nil {
 		totalRequests = -1
 	}
+	trace.Span.AddEvent(fmt.Sprintf("Total requests: %d", totalRequests))
 
 	currentContainers := 0
 	state, err := i.State()
 	if err != nil {
 		currentContainers = -1
 	}
+	trace.Span.AddEvent(fmt.Sprintf("Current containers: %d", currentContainers))
 
 	currentContainers = state.PendingContainers + state.RunningContainers