feat(gateway): report provider lifecycle

gateway/transport.go

@@ -44,18 +44,23 @@ func (f AdmissionFunc) Admit(ctx context.Context, request protocol.TunnelAdmissi
 
 func (AdmissionFunc) Release(context.Context, protocol.SessionAuthority) error { return nil }
 
+type ProviderStateReporter interface {
+	ReportProviderState(context.Context, protocol.ProviderState) error
+}
+
 type ServerConfig struct {
-	ListenAddress        string
-	TLSConfig            *tls.Config
-	QUICConfig           *quic.Config
-	GatewayID            string
-	Capabilities         protocol.CapabilityProfile
-	ProviderCapabilities protocol.CapabilityProfile
-	Admission            Admission
-	Provider             Provider
-	ProviderProfile      string
-	ProviderIdentity     string
-	PacerKbps            int64
+	ListenAddress         string
+	TLSConfig             *tls.Config
+	QUICConfig            *quic.Config
+	GatewayID             string
+	Capabilities          protocol.CapabilityProfile
+	ProviderCapabilities  protocol.CapabilityProfile
+	Admission             Admission
+	ProviderStateReporter ProviderStateReporter
+	Provider              Provider
+	ProviderProfile       string
+	ProviderIdentity      string
+	PacerKbps             int64
 }
 
 type Server struct {
@@ -212,13 +217,26 @@ func (s *Server) handleConnection(parent context.Context, connection *quic.Conn)
 		_ = writeStableError(stream, "no_capability_overlap", err, false)
 		return
 	}
+	if err := s.reportProviderState(ctx, protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateStarting, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}}); err != nil {
+		_ = s.config.Admission.Release(context.Background(), authority)
+		_ = writeStableError(stream, "provider_state_unavailable", err, true)
+		return
+	}
 	providerSession, err := s.config.Provider.Start(ctx, LaunchRequest{SessionID: request.SessionID, Capabilities: selected, ProviderProfile: authority.ProviderProfile, ProviderIdentity: authority.ProviderIdentity})
 	if err != nil {
 		s.metrics.ProviderErrors.Add(1)
+		_ = s.reportProviderState(context.Background(), protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateFailed, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}})
 		_ = s.config.Admission.Release(context.Background(), authority)
 		_ = writeStableError(stream, stableProviderCode(err), err, errors.Is(err, context.DeadlineExceeded))
 		return
 	}
+	if err := s.reportProviderState(ctx, providerSession.State()); err != nil {
+		_ = providerSession.ReleaseAll(context.Background())
+		_ = providerSession.Terminate(context.Background())
+		_ = s.config.Admission.Release(context.Background(), authority)
+		_ = writeStableError(stream, "provider_state_unavailable", err, true)
+		return
+	}
 	authority.Capabilities = selected
 	authorityBytes, err := protocol.EncodeSessionAuthority(authority)
 	if err != nil || writeWire(stream, authorityBytes, defaultHelloLimit) != nil {
@@ -237,6 +255,16 @@ func (s *Server) handleConnection(parent context.Context, connection *quic.Conn)
 	session.run()
 }
 
+func (s *Server) reportProviderState(ctx context.Context, state protocol.ProviderState) error {
+	if s.config.ProviderStateReporter == nil {
+		return nil
+	}
+	if err := state.Validate(); err != nil {
+		return err
+	}
+	return s.config.ProviderStateReporter.ReportProviderState(ctx, state)
+}
+
 func (s *Server) validateAuthority(authority protocol.SessionAuthority, request protocol.TunnelAdmissionRequest) error {
 	if err := authority.Validate(); err != nil {
 		return err
@@ -467,15 +495,32 @@ func (s *gatewaySession) cleanup() {
 		s.cancel()
 		cleanupCtx, cancel := context.WithTimeout(context.Background(), time.Second)
 		defer cancel()
-		if err := s.provider.ReleaseAll(cleanupCtx); err != nil {
+		releaseInputsErr := s.provider.ReleaseAll(cleanupCtx)
+		if releaseInputsErr != nil {
 			s.server.metrics.ProviderErrors.Add(1)
 		}
-		if err := s.provider.Terminate(cleanupCtx); err != nil {
+		terminateErr := s.provider.Terminate(cleanupCtx)
+		if terminateErr != nil {
 			s.server.metrics.ProviderErrors.Add(1)
 		}
+		state := s.provider.State()
+		if releaseInputsErr != nil || terminateErr != nil {
+			state.State = ProviderStateCleanup
+			state.CleanupPending = true
+			if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
+				s.server.metrics.ProviderErrors.Add(1)
+			} else if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
+				s.server.metrics.ProviderErrors.Add(1)
+			}
+			_ = s.connection.CloseWithError(applicationError, "session closed")
+			return
+		}
 		if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
 			s.server.metrics.ProviderErrors.Add(1)
 		}
+		if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
+			s.server.metrics.ProviderErrors.Add(1)
+		}
 		_ = s.connection.CloseWithError(applicationError, "session closed")
 	})
 }