fix

Author: maxdml

dbos/dbos.go

@@ -195,9 +195,7 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 	// Will do nothing if the concrete type is not dbosContext
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
 		launched := dbosCtx.launched.Load()
-		childCtxLaunched := atomic.Bool{}
-		childCtxLaunched.Store(launched)
-		return &dbosContext{
+		childCtx := &dbosContext{
 			ctx:                     context.WithValue(dbosCtx.ctx, key, val), // Spawn a new child context with the value set
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -207,8 +205,9 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
-			launched:                childCtxLaunched,
 		}
+		childCtx.launched.Store(launched)
+		return childCtx
 	}
 	return nil
 }
@@ -222,11 +221,9 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 	}
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
 		launched := dbosCtx.launched.Load()
-		childCtxLaunched := atomic.Bool{}
-		childCtxLaunched.Store(launched)
 		// Create a new context that is not canceled when the parent is canceled
 		// but retains all other values
-		return &dbosContext{
+		childCtx := &dbosContext{
 			ctx:                     context.WithoutCancel(dbosCtx.ctx),
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -236,8 +233,9 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
-			launched:                childCtxLaunched,
 		}
+		childCtx.launched.Store(launched)
+		return childCtx
 	}
 	return nil
 }
@@ -251,10 +249,8 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 	}
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
 		launched := dbosCtx.launched.Load()
-		childCtxLaunched := atomic.Bool{}
-		childCtxLaunched.Store(launched)
 		newCtx, cancelFunc := context.WithTimeoutCause(dbosCtx.ctx, timeout, errors.New("DBOS context timeout"))
-		return &dbosContext{
+		childCtx := &dbosContext{
 			ctx:                     newCtx,
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -264,8 +260,9 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
-			launched:                childCtxLaunched,
-		}, cancelFunc
+		}
+		childCtx.launched.Store(launched)
+		return childCtx, cancelFunc
 	}
 	return nil, func() {}
 }