the shutdown function may be nil (#353)

Author: kans

pkg/cli/commands.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	"github.com/spf13/cobra"
@@ -26,6 +27,10 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/uotel"
 )
 
+const (
+	otelShutdownTimeout = 5 * time.Second
+)
+
 type ContrainstSetter func(*cobra.Command, field.Configuration) error
 
 func MakeMainCommand[T field.Configurable](
@@ -64,7 +69,9 @@ func MakeMainCommand[T field.Configurable](
 			if otelShutdown == nil {
 				return
 			}
-			err := otelShutdown(context.Background())
+			shutdownCtx, cancel := context.WithDeadline(context.Background(), time.Now().Add(otelShutdownTimeout))
+			defer cancel()
+			err := otelShutdown(shutdownCtx)
 			if err != nil {
 				zap.L().Error("error shutting down otel", zap.Error(err))
 			}
@@ -332,7 +339,12 @@ func MakeGRPCServerCommand[T field.Configurable](
 			return err
 		}
 		defer func() {
-			err := otelShutdown(context.Background())
+			if otelShutdown == nil {
+				return
+			}
+			shutdownCtx, cancel := context.WithDeadline(context.Background(), time.Now().Add(otelShutdownTimeout))
+			defer cancel()
+			err := otelShutdown(shutdownCtx)
 			if err != nil {
 				zap.L().Error("error shutting down otel", zap.Error(err))
 			}

pkg/cli/lambda_server__added.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"time"
 
 	aws_lambda "github.com/aws/aws-lambda-go/lambda"
 	"github.com/conductorone/baton-sdk/pkg/crypto/providers/jwk"
@@ -78,7 +79,12 @@ func OptionallyAddLambdaCommand[T field.Configurable](
 			return err
 		}
 		defer func() {
-			err := otelShutdown(context.Background())
+			if otelShutdown == nil {
+				return
+			}
+			shutdownCtx, cancel := context.WithDeadline(context.Background(), time.Now().Add(otelShutdownTimeout))
+			defer cancel()
+			err := otelShutdown(shutdownCtx)
 			if err != nil {
 				zap.L().Error("error shutting down otel", zap.Error(err))
 			}
@@ -178,5 +184,4 @@ func OptionallyAddLambdaCommand[T field.Configurable](
 		return nil
 	}
 	return nil
-
 }