conn established (success) or done (end/failure..) are signals of the supervisor

2026-09-17 20:29:07 +02:00 · 2026-06-18 14:40:03 +02:00
parent 0e8fd22f36
commit 29ee84999c
4 changed files with 102 additions and 64 deletions
@@ -95,11 +95,22 @@ func (c *ConnectClient) Run(config *profilemanager.Config, md metadata.MD, runni
 	return c.sup.start(config, md, MobileDependency{}, runningChan, logPath)
 }

-// RunAsync starts a client run without blocking. done (if non-nil, buffered)
-// receives the run's end result. Used by the daemon, which drives the lifecycle
-// through the supervisor (Run/Stop/IsRunning) rather than blocking on Run.
-func (c *ConnectClient) RunAsync(config *profilemanager.Config, md metadata.MD, runningChan chan struct{}, done chan error) {
-	c.sup.startAsync(config, md, MobileDependency{}, runningChan, "", done)
+// RunAsync starts a client run without blocking. Used by the daemon, which
+// drives the lifecycle through the supervisor rather than blocking on Run; it
+// then waits for the outcome via WaitEstablishedOrDone. The run's lifecycle
+// channels are created and owned by the supervisor — callers never hold them.
+func (c *ConnectClient) RunAsync(config *profilemanager.Config, md metadata.MD) {
+	est := make(chan struct{})
+	d := make(chan error, 1)
+	c.sup.startAsync(config, md, MobileDependency{}, est, "", d)
+}
+
+// WaitEstablishedOrDone blocks until the in-flight run becomes established (nil),
+// ends before that (the run error, or a sentinel on a clean stop), or ctx is
+// cancelled. Returns errNoRunInFlight if no run is in flight. Wraps the wait on
+// the supervisor-owned channels so callers never touch them directly.
+func (c *ConnectClient) WaitEstablishedOrDone(ctx context.Context) error {
+	return c.sup.waitEstablishedOrDone(ctx)
 }

 // RunOnAndroid with main logic on mobile system
@@ -13,6 +13,13 @@ import (
 // in flight.
 var errAlreadyRunning = errors.New("client is already running")

+// errNoRunInFlight is returned by waitEstablishedOrDone when no run is active.
+var errNoRunInFlight = errors.New("no connection run in flight")
+
+// errStoppedBeforeEstablished is returned when a run ended (cleanly) before the
+// connection was established.
+var errStoppedBeforeEstablished = errors.New("run stopped before the connection was established")
+
 // lifecycleOp is a serialized lifecycle operation processed by the supervisor.
 type lifecycleOp int

@@ -20,6 +27,7 @@ const (
 	opStart lifecycleOp = iota
 	opStop
 	opStatus
+	opSignals
 )

 // lifecycleCmd is a single start/stop/status request handed to the supervisor
@@ -31,7 +39,7 @@ const (
 //     errAlreadyRunning immediately if a run is already in flight.
 //   - for opStop it receives nil once the in-flight run has fully unwound.
 //
-// reply is used only by opStatus: it receives whether a run is in flight.
+// reply is used only by opStatus; sigReply only by opSignals.
 type lifecycleCmd struct {
 	op          lifecycleOp
 	config      *profilemanager.Config
@@ -41,6 +49,14 @@ type lifecycleCmd struct {
 	logPath     string
 	done        chan error
 	reply       chan bool
+	sigReply    chan runSignals
+}
+
+// runSignals exposes the in-flight run's lifecycle channels to external waiters
+// (the daemon's waitForUp/Status). Both are nil when no run is in flight.
+type runSignals struct {
+	established <-chan struct{} // closed by the run once the connection is established
+	done        <-chan error    // receives the run's end result
 }

 // runEndResult is sent by the run goroutine to the supervisor when a run ends,
@@ -96,6 +112,12 @@ func (s *supervisor) loop() {
 				s.handleStop(cmd)
 			case opStatus:
 				cmd.reply <- (s.isRunningInternal())
+			case opSignals:
+				var sig runSignals
+				if s.curStart != nil {
+					sig = runSignals{established: s.curStart.runningChan, done: s.curStart.done}
+				}
+				cmd.sigReply <- sig
 			}
 		case res := <-s.runEnded:
 			// Run ended on its own, without an explicit Stop.
@@ -216,6 +238,46 @@ func (s *supervisor) isRunningInternal() bool {
 	return s.curStart != nil
 }

+// waitEstablishedOrDone blocks until the in-flight run becomes established
+// (returns nil) or ends before that (returns the run error, or
+// errStoppedBeforeEstablished on a clean stop), or ctx is cancelled. Returns
+// errNoRunInFlight if no run is in flight. The select runs in the caller's
+// goroutine on the run's channels — it does not block the supervisor loop.
+func (s *supervisor) waitEstablishedOrDone(ctx context.Context) error {
+	sig := s.signals()
+	if sig.established == nil {
+		return errNoRunInFlight
+	}
+	select {
+	case <-sig.established:
+		return nil
+	case err := <-sig.done:
+		if err != nil {
+			return err
+		}
+		return errStoppedBeforeEstablished
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// signals asks the loop for the in-flight run's lifecycle channels, serialized
+// with start/stop so the returned pair is consistent. Both nil when idle.
+func (s *supervisor) signals() runSignals {
+	reply := make(chan runSignals, 1)
+	select {
+	case s.cmdCh <- lifecycleCmd{op: opSignals, sigReply: reply}:
+	case <-s.ctx.Done():
+		return runSignals{}
+	}
+	select {
+	case sig := <-reply:
+		return sig
+	case <-s.ctx.Done():
+		return runSignals{}
+	}
+}
+
 // stop enqueues a stop and blocks until the in-flight run is fully torn down.
 func (s *supervisor) stop() error {
 	done := make(chan error, 1)
@@ -142,12 +142,10 @@ func (s *Server) restartEngineForMDMLocked() error {

 	_, cancel := context.WithCancel(s.rootCtx)
 	s.actCancel = cancel
-	s.connectionEstablishedChan = make(chan struct{})
-	s.connectionDoneChan = make(chan error, 1)
 	log.Info("MDM restart: starting a fresh run with re-resolved config")
-	// MDM restart has no incoming RPC metadata; fire and forget (the supervisor
-	// reconnects internally and we don't block on the run).
-	s.connectClient.RunAsync(config, nil, s.connectionEstablishedChan, s.connectionDoneChan)
+	// MDM restart has no incoming RPC metadata; fire and forget (the run owns
+	// its established/done channels, the supervisor reconnects internally).
+	s.connectClient.RunAsync(config, nil)
 	s.publishConfigChangedEvent("mdm")
 	return nil
 }
@@ -62,10 +62,8 @@ type Server struct {
 	mutex  sync.Mutex
 	config *profilemanager.Config
 	proto.UnimplementedDaemonServiceServer
-	// Whether a run is in flight is owned by the supervisor
-	// (connectClient.ConnectionRunning); the daemon keeps no separate flag.
-	connectionEstablishedChan chan struct{} // closed by the run once the connection is established (StatusConnected)
-	connectionDoneChan        chan error    // receives the run's end result
+	// Run state (in-flight? established/done channels?) is owned entirely by the
+	// supervisor inside connectClient — the daemon keeps no per-run fields.

 	connectClient *internal.ConnectClient

@@ -227,11 +225,10 @@ func (s *Server) Start() error {
 		return nil
 	}

-	s.connectionEstablishedChan = make(chan struct{})
-	s.connectionDoneChan = make(chan error, 1)
 	// Boot autoconnect: no incoming RPC metadata. The supervisor runs the
-	// client and reconnects internally; we just fire and forget.
-	s.connectClient.RunAsync(config, nil, s.connectionEstablishedChan, s.connectionDoneChan)
+	// client and reconnects internally; we just fire and forget (the run owns
+	// its established/done channels).
+	s.connectClient.RunAsync(config, nil)
 	s.publishConfigChangedEvent("startup")
 	return nil
 }
@@ -753,47 +750,26 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
 	s.statusRecorder.UpdateManagementAddress(s.config.ManagementURL.String())
 	s.statusRecorder.UpdateRosenpass(s.config.RosenpassEnabled, s.config.RosenpassPermissive)

-	s.connectionEstablishedChan = make(chan struct{})
-	s.connectionDoneChan = make(chan error, 1)
-
-	s.connectClient.RunAsync(s.config, md, s.connectionEstablishedChan, s.connectionDoneChan)
+	s.connectClient.RunAsync(s.config, md)
 	s.publishConfigChangedEvent("up_rpc")

 	s.mutex.Unlock()
 	return s.waitForUp(callerCtx)
 }

+// waitForUp blocks until the in-flight run becomes established (success) or ends
+// before that (failure). The wait is owned by the supervisor (via the client) —
+// the daemon holds no per-run state here.
 func (s *Server) waitForUp(callerCtx context.Context) (*proto.UpResponse, error) {
 	timeoutCtx, cancel := context.WithTimeout(callerCtx, 50*time.Second)
 	defer cancel()

-	// Read the per-run channels under the lock. They are written under s.mutex
-	// by Up/Start and by the MDM restart (restartEngineForMDMLocked, which runs
-	// on the ticker goroutine), so reading them here — where the Up caller has
-	// already released the lock — must be synchronized both to avoid a data race
-	// and to capture a consistent (established, done) pair from the same run.
-	s.mutex.Lock()
-	establishedChan := s.connectionEstablishedChan
-	doneChan := s.connectionDoneChan
-	s.mutex.Unlock()
-
-	select {
-	case err := <-doneChan:
-		// The run ended before signalling ready: it failed or was stopped.
-		if err != nil {
-			return nil, fmt.Errorf("client failed to connect: %w", err)
-		}
-		return nil, fmt.Errorf("client stopped before becoming ready")
-	case <-establishedChan:
-		s.isSessionActive.Store(true)
-		return &proto.UpResponse{}, nil
-	case <-callerCtx.Done():
-		log.Debug("context done, stopping the wait for engine to become ready")
-		return nil, callerCtx.Err()
-	case <-timeoutCtx.Done():
-		log.Debug("up is timed out, stopping the wait for engine to become ready")
-		return nil, timeoutCtx.Err()
+	if err := s.connectClient.WaitEstablishedOrDone(timeoutCtx); err != nil {
+		log.Debugf("waiting for the connection to be established failed: %v", err)
+		return nil, fmt.Errorf("connection not established: %w", err)
 	}
+	s.isSessionActive.Store(true)
+	return &proto.UpResponse{}, nil
 }

 func (s *Server) switchProfileIfNeeded(profileName string, userName *string, activeProf *profilemanager.ActiveProfileState) error {
@@ -1078,21 +1054,12 @@ func (s *Server) Status(
 	ctx context.Context,
 	msg *proto.StatusRequest,
 ) (*proto.StatusResponse, error) {
-	// Snapshot the run state under the lock. A run that hits a terminal auth
-	// failure now exits on its own (engine marks NeedsLogin), so we no longer
-	// poll-and-cancel: we just wait for the run to become ready or to end.
-	s.mutex.Lock()
-	client := s.connectClient
-	establishedChan := s.connectionEstablishedChan
-	doneChan := s.connectionDoneChan
-	s.mutex.Unlock()
-	alive := client.ConnectionRunning()
-
-	if msg.WaitForReady != nil && *msg.WaitForReady && alive {
-		select {
-		case <-establishedChan:
-		case <-doneChan:
-		case <-ctx.Done():
+	// A run that hits a terminal auth failure now exits on its own (engine marks
+	// NeedsLogin), so we no longer poll-and-cancel: we wait for the in-flight run
+	// to become established or to end. With no run in flight this returns
+	// immediately (errNoRunInFlight); either way we then report the status below.
+	if msg.WaitForReady != nil && *msg.WaitForReady {
+		if err := s.connectClient.WaitEstablishedOrDone(ctx); err != nil && ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
 	}