HH-690: hold data lock through warm startup

Co-authored-by: multica-agent <github@multica.ai>
2026-08-26 15:56:09 +08:00
parent e1136c6a0b
commit 175c309fcb
2 changed files with 175 additions and 75 deletions
@@ -38,10 +38,11 @@ type CandidateConfig struct {
 }

 type LifecycleConfig struct {
-	Candidate      CandidateConfig
-	ControllerURL  string
-	UpdateInterval time.Duration
-	Trigger        <-chan os.Signal
+	Candidate              CandidateConfig
+	ControllerURL          string
+	UpdateInterval         time.Duration
+	Trigger                <-chan os.Signal
+	afterLastGoodValidated func()
 }

 // PublishCandidate performs one Stage 1 update. Starting or reloading Mihomo is
@@ -155,63 +156,83 @@ func Run(ctx context.Context, config LifecycleConfig) error {
 		return err
 	}

-	var generation string
-	valid := withDataDirLock(config.Candidate.DataDir, func(dataDir string) error {
-		config.Candidate.DataDir = dataDir
-		var err error
-		generation, err = validLastGoodLocked(ctx, config.Candidate)
-		return err
-	})
-	warm := valid == nil
-	if !warm {
-		if err := PublishCandidate(ctx, config.Candidate); err != nil {
-			return fmt.Errorf("cold-start candidate failed: %w", err)
-		}
-		valid = withDataDirLock(config.Candidate.DataDir, func(dataDir string) error {
-			config.Candidate.DataDir = dataDir
-			var err error
-			generation, err = validLastGoodLocked(ctx, config.Candidate)
-			return err
-		})
-		if valid != nil {
-			return fmt.Errorf("published candidate is invalid: %w", valid)
-		}
-	}
-
-	serviceCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	mihomo := serviceCommand(serviceCtx, config.Candidate.MihomoBinary, generation, filepath.Join(generation, "config.yaml"))
-	if err := mihomo.Start(); err != nil {
-		return fmt.Errorf("start Mihomo: %w", err)
-	}
-	exit := make(chan error, 1)
-	go func() { exit <- mihomo.Wait() }()
 	client := config.Candidate.Client
 	if client == nil {
 		client = &http.Client{Timeout: 30 * time.Second}
 	}
-	if err := waitForController(ctx, client, controllerURL, exit); err != nil {
-		return err
+	var cancel context.CancelFunc
+	var exit chan error
+	var done chan struct{}
+	startupErr := withDataDirLock(config.Candidate.DataDir, func(dataDir string) error {
+		config.Candidate.DataDir = dataDir
+		generation, valid := validLastGoodLocked(ctx, config.Candidate)
+		warm := valid == nil
+		if !warm {
+			if err := publishCandidateLocked(ctx, config.Candidate); err != nil {
+				return fmt.Errorf("cold-start candidate failed: %w", err)
+			}
+			generation, valid = validLastGoodLocked(ctx, config.Candidate)
+			if valid != nil {
+				return fmt.Errorf("published candidate is invalid: %w", valid)
+			}
+		}
+		if warm && config.afterLastGoodValidated != nil {
+			config.afterLastGoodValidated()
+		}
+
+		serviceCtx, serviceCancel := context.WithCancel(ctx)
+		cancel = serviceCancel
+		mihomo := serviceCommand(serviceCtx, config.Candidate.MihomoBinary, generation, filepath.Join(generation, "config.yaml"))
+		if err := mihomo.Start(); err != nil {
+			cancel()
+			return fmt.Errorf("start Mihomo: %w", err)
+		}
+		exit = make(chan error, 1)
+		done = make(chan struct{})
+		go func() {
+			exit <- mihomo.Wait()
+			close(done)
+		}()
+		stop := func() {
+			cancel()
+			<-done
+		}
+		if err := waitForController(ctx, client, controllerURL, exit); err != nil {
+			stop()
+			return err
+		}
+		log.Printf("bootstrap: Mihomo started config=%s update_interval=%s", filepath.Base(generation), config.UpdateInterval)
+		if warm {
+			if err := updateAndReloadLocked(ctx, client, controllerURL, config.Candidate); err != nil {
+				stop()
+				return fmt.Errorf("fatal update stopped Mihomo: %w", err)
+			}
+		}
+		return nil
+	})
+	if startupErr != nil {
+		if cancel != nil {
+			cancel()
+			if done != nil {
+				<-done
+			}
+		}
+		return startupErr
 	}
-	log.Printf("bootstrap: Mihomo started config=%s update_interval=%s", filepath.Base(generation), config.UpdateInterval)
+	defer cancel()
 	failClosed := func(err error) error {
 		cancel()
-		<-exit
+		<-done
 		return fmt.Errorf("fatal update stopped Mihomo: %w", err)
 	}

-	if warm {
-		if err := updateAndReload(ctx, client, controllerURL, config.Candidate); err != nil {
-			return failClosed(err)
-		}
-	}
 	ticker := time.NewTicker(config.UpdateInterval)
 	defer ticker.Stop()
 	for {
 		select {
 		case <-ctx.Done():
 			cancel()
-			<-exit
+			<-done
 			return ctx.Err()
 		case err := <-exit:
 			if err == nil {
@@ -256,39 +277,44 @@ func validLastGoodLocked(ctx context.Context, config CandidateConfig) (string, e
 func updateAndReload(ctx context.Context, client *http.Client, controllerURL string, config CandidateConfig) error {
 	return withDataDirLock(config.DataDir, func(dataDir string) error {
 		config.DataDir = dataDir
-		previous, err := currentGeneration(filepath.Join(dataDir, "last-good"))
-		if err != nil || previous == "" {
-			return errors.Join(errors.New("last-good is unavailable during update"), err)
-		}
-		if err := publishCandidateLocked(ctx, config); err != nil {
-			if errors.Is(err, errPersistence) {
-				return err
-			}
-			log.Printf("bootstrap: update rejected; keeping last-good: %v", err)
-			return nil
-		}
-		next, err := currentGeneration(filepath.Join(dataDir, "last-good"))
-		if err == nil {
-			err = reloadMihomo(ctx, client, controllerURL, filepath.Join(dataDir, filepath.FromSlash(next), "config.yaml"))
-		}
-		if err == nil {
-			log.Print("bootstrap: configuration updated and reloaded")
-			return nil
-		}
-		if rollbackErr := config.replacePointer(filepath.Join(dataDir, "last-good"), previous); rollbackErr != nil {
-			return fmt.Errorf("restore last-good pointer after reload failure: %w", rollbackErr)
-		}
-		if rollbackErr := config.syncDirectory(dataDir); rollbackErr != nil {
-			return fmt.Errorf("persist restored last-good pointer after reload failure: %w", rollbackErr)
-		}
-		if rollbackErr := reloadMihomo(ctx, client, controllerURL, filepath.Join(dataDir, filepath.FromSlash(previous), "config.yaml")); rollbackErr != nil {
-			return fmt.Errorf("reload restored last-good after reload failure: %w", rollbackErr)
-		}
-		log.Printf("bootstrap: reload rejected; restored and reloaded last-good: %v", err)
-		return nil
+		return updateAndReloadLocked(ctx, client, controllerURL, config)
 	})
 }

+func updateAndReloadLocked(ctx context.Context, client *http.Client, controllerURL string, config CandidateConfig) error {
+	dataDir := config.DataDir
+	previous, err := currentGeneration(filepath.Join(dataDir, "last-good"))
+	if err != nil || previous == "" {
+		return errors.Join(errors.New("last-good is unavailable during update"), err)
+	}
+	if err := publishCandidateLocked(ctx, config); err != nil {
+		if errors.Is(err, errPersistence) {
+			return err
+		}
+		log.Printf("bootstrap: update rejected; keeping last-good: %v", err)
+		return nil
+	}
+	next, err := currentGeneration(filepath.Join(dataDir, "last-good"))
+	if err == nil {
+		err = reloadMihomo(ctx, client, controllerURL, filepath.Join(dataDir, filepath.FromSlash(next), "config.yaml"))
+	}
+	if err == nil {
+		log.Print("bootstrap: configuration updated and reloaded")
+		return nil
+	}
+	if rollbackErr := config.replacePointer(filepath.Join(dataDir, "last-good"), previous); rollbackErr != nil {
+		return fmt.Errorf("restore last-good pointer after reload failure: %w", rollbackErr)
+	}
+	if rollbackErr := config.syncDirectory(dataDir); rollbackErr != nil {
+		return fmt.Errorf("persist restored last-good pointer after reload failure: %w", rollbackErr)
+	}
+	if rollbackErr := reloadMihomo(ctx, client, controllerURL, filepath.Join(dataDir, filepath.FromSlash(previous), "config.yaml")); rollbackErr != nil {
+		return fmt.Errorf("reload restored last-good after reload failure: %w", rollbackErr)
+	}
+	log.Printf("bootstrap: reload rejected; restored and reloaded last-good: %v", err)
+	return nil
+}
+
 func reloadMihomo(ctx context.Context, client *http.Client, controllerURL, configPath string) error {
 	body, err := json.Marshal(map[string]string{"path": configPath})
 	if err != nil {
@@ -137,6 +137,80 @@ func TestRunAndCandidateSerializeDataDirUpdates(t *testing.T) {
 	}
 }

+func TestRunWarmStartKeepsValidatedGenerationLocked(t *testing.T) {
+	fixture := newLifecycleFixture(t)
+	if err := PublishCandidate(context.Background(), fixture.config); err != nil {
+		t.Fatal(err)
+	}
+	firstTarget := readLastGood(t, fixture.config.DataDir)
+	fixture.setSubscription("second-node", http.StatusOK)
+	validated := make(chan struct{})
+	resume := make(chan struct{})
+	lifecycle := fixture.lifecycle(nil)
+	lifecycle.afterLastGoodValidated = func() {
+		close(validated)
+		<-resume
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan error, 1)
+	go func() { result <- Run(ctx, lifecycle) }()
+	select {
+	case <-validated:
+	case <-time.After(5 * time.Second):
+		close(resume)
+		t.Fatal("Run did not pause after warm generation validation")
+	}
+
+	candidateStarted := make(chan struct{})
+	candidateResult := make(chan error, 1)
+	go func() {
+		close(candidateStarted)
+		candidateResult <- PublishCandidate(context.Background(), fixture.config)
+	}()
+	<-candidateStarted
+	select {
+	case err := <-candidateResult:
+		close(resume)
+		t.Fatalf("candidate bypassed the warm-start data lock: %v", err)
+	case <-time.After(100 * time.Millisecond):
+	}
+	if got := fixture.subscriptionRequests.Load(); got != 1 {
+		close(resume)
+		t.Fatalf("subscription requests = %d before warm start resumed, want 1", got)
+	}
+	if _, err := os.Stat(fixture.startedPath); !errors.Is(err, os.ErrNotExist) {
+		close(resume)
+		t.Fatalf("Mihomo started before warm validation resumed: %v", err)
+	}
+	close(resume)
+	fixture.waitStarted(t)
+	if err := <-candidateResult; err != nil {
+		t.Fatalf("concurrent PublishCandidate() error = %v", err)
+	}
+
+	activeTarget := "generations/a"
+	if firstTarget == activeTarget {
+		activeTarget = "generations/b"
+	}
+	if _, err := os.Stat(filepath.Join(fixture.config.DataDir, filepath.FromSlash(activeTarget))); err != nil {
+		t.Fatalf("active generation %q was removed: %v", activeTarget, err)
+	}
+	if _, err := os.Stat(filepath.Join(fixture.config.DataDir, "last-good")); err != nil {
+		t.Fatalf("last-good is dangling: %v", err)
+	}
+	select {
+	case err := <-result:
+		t.Fatalf("Run() stopped after concurrent candidate: %v", err)
+	default:
+	}
+
+	cancel()
+	if err := <-result; !errors.Is(err, context.Canceled) {
+		t.Fatalf("Run() error = %v, want context cancellation", err)
+	}
+}
+
 func TestRunRollbackPersistenceFailureStopsMihomo(t *testing.T) {
 	fixture := newLifecycleFixture(t)
 	if err := PublishCandidate(context.Background(), fixture.config); err != nil {