allow startup to be cancelled

beacon/beacon.go

@@ -72,6 +72,9 @@ type Beacon struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 	wg     sync.WaitGroup
+
+	selfRegistered bool
+	stopOnce       sync.Once
 }
 
 // New creates a beacon. Call Start to begin discovery and membership.
@@ -197,12 +200,12 @@ func (b *Beacon) applyQueuedRemovals() {
 // remote effects that arrive before its first subscription.
 func (b *Beacon) Start(ctx context.Context) (err error) {
 	b.ctx, b.cancel = context.WithCancel(ctx)
-	// A failed Start must not leak background workers: the caller only pairs
-	// Stop with a successful Start.
+	// A failed Start owns the same one-shot cleanup as normal shutdown. This
+	// rolls back a completed self-registration before returning the error and
+	// prevents a later Runtime.Stop from removing it twice.
 	defer func() {
 		if err != nil {
-			b.cancel()
-			b.wg.Wait()
+			b.Stop()
 		}
 	}()
 
@@ -273,22 +276,16 @@ func (b *Beacon) Start(ctx context.Context) (err error) {
 
 // Stop sends a REMOVE_OP for this node and stops background loops.
 func (b *Beacon) Stop() {
-	// Graceful departure: remove our membership entry. A failure here means
-	// peers keep our entry and redial us until an operator intervenes, so it
-	// must be loud.
-	ctx := b.engine.NewContext()
-	if err := ctx.Emit(buildMemberRemove(uint64(b.cfg.NodeID))); err != nil {
-		slog.Error("beacon: departure remove emit failed", "node_id", b.cfg.NodeID, "error", err)
-	} else if err := ctx.Flush(); err != nil {
-		slog.Error("beacon: departure remove flush failed", "node_id", b.cfg.NodeID, "error", err)
-	}
+	b.stopOnce.Do(func() {
+		b.unregisterSelf()
 
-	if b.cancel != nil {
-		b.cancel()
-	}
-	b.wg.Wait()
+		if b.cancel != nil {
+			b.cancel()
+		}
+		b.wg.Wait()
 
-	slog.Info("beacon stopped", "node_id", b.cfg.NodeID)
+		slog.Info("beacon stopped", "node_id", b.cfg.NodeID)
+	})
 }
 
 // Members returns the current known membership.
@@ -338,14 +335,44 @@ func (b *Beacon) primeSubscription() {
 // the stale entry becomes visible — a one-shot sweep at registration would
 // silently miss an entry that hasn't bootstrapped in yet.
 func (b *Beacon) registerSelf() error {
+	if b.selfRegistered {
+		return nil
+	}
+
 	ctx := b.engine.NewContext()
 	if err := ctx.Emit(buildMemberInsert(uint64(b.cfg.NodeID), b.cfg.AdvertiseAddr)); err != nil {
 		return err
 	}
 	if err := ctx.Emit(buildMemberTypeTag()); err != nil {
 		return err
 	}
-	return ctx.Flush()
+	if err := ctx.Flush(); err != nil {
+		return err
+	}
+	b.selfRegistered = true
+	return nil
+}
+
+// unregisterSelf removes this node's membership entry after a successful
+// registration. Stop's sync.Once makes this shared normal-shutdown and
+// failed-start rollback path one-shot.
+func (b *Beacon) unregisterSelf() {
+	if !b.selfRegistered {
+		return
+	}
+
+	// A failure here means peers keep our entry and redial us until an operator
+	// intervenes, so it must be loud.
+	ctx := b.engine.NewContext()
+	if err := ctx.Emit(buildMemberRemove(uint64(b.cfg.NodeID))); err != nil {
+		slog.Error("beacon: departure remove emit failed", "node_id", b.cfg.NodeID, "error", err)
+		return
+	}
+	if err := ctx.Flush(); err != nil {
+		slog.Error("beacon: departure remove flush failed", "node_id", b.cfg.NodeID, "error", err)
+		return
+	}
+	b.selfRegistered = false
 }
 
 // topologyLoop keeps the PeerManager's connection table as a reactive

beacon/runtime.go

@@ -147,7 +147,7 @@ type Runtime struct {
 // adds its own callbacks via Engine.OnKeyDataAdded etc. is
 // responsible for unwiring those — Stop only tears down what
 // NewRuntime built.
-func NewRuntime(cfg RuntimeConfig) (*Runtime, error) {
+func NewRuntime(ctx context.Context, cfg RuntimeConfig) (*Runtime, error) {
 	log := cfg.Logger
 	if log == nil {
 		log = slog.Default()
@@ -261,7 +261,7 @@ func NewRuntime(cfg RuntimeConfig) (*Runtime, error) {
 	// connection establishment, and a peer counts as "holds everything"
 	// until its filter arrives.
 	pm.SetKeyFilterProvider(engine)
-	if err := pm.Start(context.Background()); err != nil {
+	if err := pm.Start(ctx); err != nil {
 		return nil, errors.Join(
 			fmt.Errorf("failed to start peer manager: %w", err),
 			closeRuntimeOnError(rt),
@@ -333,7 +333,7 @@ func NewRuntime(cfg RuntimeConfig) (*Runtime, error) {
 		// promptly and serve solo until peers converge. joinCtx lets Stop
 		// abort an in-flight join; joinWG lets it wait for the goroutine to
 		// return before closing the peer manager / engine the join uses.
-		joinCtx, joinCancel := context.WithCancel(context.Background())
+		joinCtx, joinCancel := context.WithCancel(ctx)
 		rt.joinCancel = joinCancel
 		rt.Beacon = b
 		rt.joinWG.Go(func() {
@@ -342,7 +342,7 @@ func NewRuntime(cfg RuntimeConfig) (*Runtime, error) {
 			}
 		})
 	} else {
-		if err := b.Start(context.Background()); err != nil {
+		if err := b.Start(ctx); err != nil {
 			return nil, errors.Join(
 				fmt.Errorf("failed to start beacon: %w", err),
 				closeRuntimeOnError(rt),

readme.md

@@ -281,7 +281,7 @@ Most of the time you don't touch `Beacon` directly — you build a `Runtime`, wh
 whole engine:
 
 ```go
-rt, err := beacon.NewRuntime(beacon.RuntimeConfig{
+rt, err := beacon.NewRuntime(context.Background(), beacon.RuntimeConfig{
     ClusterPassphrase: passphrase, // empty → single-node, no peers, no beacon
     ClusterPort:       7000,
     JoinAddr:          "my-cache.local", // DNS name every peer resolves to