pass around launched when creating child contexts

Author: maxdml

dbos/dbos.go

@@ -194,6 +194,9 @@ 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{
 			ctx:                     context.WithValue(dbosCtx.ctx, key, val), // Spawn a new child context with the value set
 			logger:                  dbosCtx.logger,
@@ -204,6 +207,7 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
+			launched:                childCtxLaunched,
 		}
 	}
 	return nil
@@ -217,6 +221,11 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 		return nil
 	}
 	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{
 			ctx:                     context.WithoutCancel(dbosCtx.ctx),
 			logger:                  dbosCtx.logger,
@@ -227,6 +236,7 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
+			launched:                childCtxLaunched,
 		}
 	}
 	return nil
@@ -240,6 +250,9 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 		return nil, func() {}
 	}
 	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{
 			ctx:                     newCtx,
@@ -251,6 +264,7 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
+			launched:                childCtxLaunched,
 		}, cancelFunc
 	}
 	return nil, func() {}