shared(api): context infrastructure plumbing for shared packages (#248)

packages/api/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
@@ -121,6 +122,9 @@ func NewGinServer(apiStore *handlers.APIStore, swagger *openapi3.T, port int) *h
 }
 
 func main() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	fmt.Println("Initializing...")
 
 	port := flag.Int("port", defaultPort, "Port for test HTTP server")
@@ -139,8 +143,8 @@ func main() {
 	}
 
 	if !env.IsLocal() {
-		shutdown := telemetry.InitOTLPExporter(serviceName, swagger.Info.Version)
-		defer shutdown()
+		shutdown := telemetry.InitOTLPExporter(ctx, serviceName, swagger.Info.Version)
+		defer shutdown(context.TODO())
 	}
 
 	// Create an instance of our handler which satisfies the generated interface

packages/orchestrator/cmd/mock-snapshot/mock.go

@@ -189,7 +189,7 @@ func mockSnapshot(
 
 	fmt.Println("Snapshotting sandbox")
 
-	snapshot, err := sbx.Snapshot(ctx, snapshotTemplateFiles, func() {})
+	snapshot, err := sbx.Snapshot(ctx, otel.Tracer("orchestrator-mock"), snapshotTemplateFiles, func() {})
 	if err != nil {
 		return fmt.Errorf("failed to snapshot sandbox: %w", err)
 	}

packages/orchestrator/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
@@ -14,14 +15,16 @@ import (
 const defaultPort = 5008
 
 func main() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	port := flag.Int("port", defaultPort, "orchestrator server port")
 
 	flag.Parse()
 
 	if !env.IsLocal() {
-		shutdown := telemetry.InitOTLPExporter(server.ServiceName, "no")
-
-		defer shutdown()
+		shutdown := telemetry.InitOTLPExporter(ctx, server.ServiceName, "no")
+		defer shutdown(context.TODO())
 	}
 
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))

packages/shared/pkg/db/client.go

@@ -36,6 +36,6 @@ func NewClient() (*DB, error) {
 	return &DB{Client: client}, nil
 }
 
-func (db *DB) Close() {
-	_ = db.Client.Close()
+func (db *DB) Close() error {
+	return db.Client.Close()
 }

packages/shared/pkg/telemetry/otel.go

@@ -33,7 +33,7 @@ type client struct {
 }
 
 // InitOTLPExporter initializes an OTLP exporter, and configures the corresponding trace providers.
-func InitOTLPExporter(serviceName, serviceVersion string) func() {
+func InitOTLPExporter(ctx context.Context, serviceName, serviceVersion string) func(ctx context.Context) error {
 	attributes := []attribute.KeyValue{
 		semconv.ServiceName(serviceName),
 		semconv.ServiceVersion(serviceVersion),
@@ -46,8 +46,6 @@ func InitOTLPExporter(serviceName, serviceVersion string) func() {
 		attributes = append(attributes, semconv.HostName(hostname))
 	}
 
-	ctx := context.Background()
-
 	res, err := resource.New(
 		ctx,
 		resource.WithSchemaURL(semconv.SchemaURL),
@@ -126,23 +124,21 @@ func InitOTLPExporter(serviceName, serviceVersion string) func() {
 	}()
 
 	// Shutdown will flush any remaining spans and shut down the exporter.
-	return func() {
-		otelClient.close()
-	}
+	return otelClient.close
 }
 
-func (c *client) close() {
-	ctx := context.Background()
-
+func (c *client) close(ctx context.Context) error {
 	if c.tracerProvider != nil {
 		if err := c.tracerProvider.Shutdown(ctx); err != nil {
-			log.Fatal(err)
+			return err
 		}
 	}
 
 	if c.meterProvider != nil {
 		if err := c.meterProvider.Shutdown(ctx); err != nil {
-			log.Fatal(err)
+			return err
 		}
 	}
+
+	return nil
 }

packages/template-manager/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
@@ -19,6 +20,9 @@ const (
 )
 
 func main() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	testFlag := flag.String("test", "", "run tests")
 	templateID := flag.String("template", "", "template id")
 	buildID := flag.String("build", "", "build id")
@@ -41,8 +45,8 @@ func main() {
 	}
 
 	if !env.IsLocal() {
-		shutdown := telemetry.InitOTLPExporter(constants.ServiceName, "no")
-		defer shutdown()
+		shutdown := telemetry.InitOTLPExporter(ctx, constants.ServiceName, "no")
+		defer shutdown(context.TODO())
 	}
 
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))