ai/live: Store and forward parameter updates. (#3682)

This allows us to correctly handle cases where:

* We haven't selected an orchestrator yet but the client sent a
  control API update.

* We've swapped orchestrators.

* We're mid-swap without an orchestrator.

In all these cases, the newly selected orchestrator will receive
the most recent parameter update as soon as it's selected.

Also:

* Add a 1 MB limit for reading parameters.

* Unblock processStream if selecting the first orchestrator fails.

Author: j0sh

core/livepeernode.go

@@ -170,6 +170,7 @@ type LivepeerNode struct {
 
 type LivePipeline struct {
 	RequestID   string
+	Params      string
 	ControlPub  *trickle.TricklePublisher
 	StopControl func()
 }

server/ai_live_video.go

@@ -452,6 +452,24 @@ func copySegment(ctx context.Context, segment *http.Response, w io.Writer, seq i
 	}
 }
 
+func registerControl(ctx context.Context, params aiRequestParams) {
+	params.node.LiveMu.Lock()
+	defer params.node.LiveMu.Unlock()
+
+	stream := params.liveParams.stream
+	if sess, exists := params.node.LivePipelines[stream]; exists && sess.ControlPub != nil {
+		if sess.ControlPub != nil {
+			clog.Info(ctx, "Stopping existing control loop", "existing_request_id", sess.RequestID)
+			sess.ControlPub.Close()
+			// TODO better solution than allowing existing streams to stomp over one another
+		}
+	}
+
+	params.node.LivePipelines[stream] = &core.LivePipeline{
+		RequestID: params.liveParams.requestID,
+	}
+}
+
 func startControlPublish(ctx context.Context, control *url.URL, params aiRequestParams) {
 	stream := params.liveParams.stream
 	controlPub, err := trickle.NewTricklePublisher(control.String())
@@ -472,22 +490,40 @@ func startControlPublish(ctx context.Context, control *url.URL, params aiRequest
 		})
 	}
 
-	if control, exists := params.node.LivePipelines[stream]; exists {
-		clog.Info(ctx, "Stopping existing control loop", "existing_request_id", control.RequestID)
-		control.ControlPub.Close()
-		// TODO better solution than allowing existing streams to stomp over one another
+	sess, exists := params.node.LivePipelines[stream]
+	if !exists || sess.RequestID != params.liveParams.requestID {
+		stopProcessing(ctx, params, fmt.Errorf("control session did not exist"))
+		return
 	}
-
-	params.node.LivePipelines[stream] = &core.LivePipeline{
-		ControlPub:  controlPub,
-		StopControl: stop,
-		RequestID:   params.liveParams.requestID,
+	if sess.ControlPub != nil {
+		// clean up from existing orchestrator
+		go sess.ControlPub.Close()
 	}
+	sess.ControlPub = controlPub
+	sess.StopControl = stop
+
 	if monitor.Enabled {
 		monitor.AICurrentLiveSessions(len(params.node.LivePipelines))
 		logCurrentLiveSessions(params.node.LivePipelines)
 	}
 
+	// Send any cached control params in a goroutine outside the lock.
+	msg := sess.Params
+	go func() {
+		if msg == "" {
+			return
+		}
+		var err error
+		for i := 0; i < 3; i++ {
+			err = controlPub.Write(strings.NewReader(msg))
+			if err == nil {
+				return
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+		stopProcessing(ctx, params, fmt.Errorf("control write failed: %w", err))
+	}()
+
 	// send a keepalive periodically to keep both ends of the connection alive
 	go func() {
 		for {

server/ai_mediaserver.go

@@ -633,6 +633,8 @@ func (ls *LivepeerServer) StartLiveVideo() http.Handler {
 			},
 		}
 
+		registerControl(ctx, params)
+
 		// Create a special parent context for orchestrator cancellation
 		orchCtx, orchCancel := context.WithCancel(ctx)
 
@@ -721,6 +723,10 @@ func processStream(ctx context.Context, params aiRequestParams, req worker.GenLi
 			}
 			params.liveParams.sendErrorEvent(err)
 		}
+		if isFirst {
+			// failed before selecting an orchestrator
+			firstProcessed <- struct{}{}
+		}
 		params.liveParams.kickInput(err)
 	}()
 	<-firstProcessed
@@ -820,22 +826,36 @@ func (ls *LivepeerServer) UpdateLiveVideo() http.Handler {
 			return
 		}
 		ls.LivepeerNode.LiveMu.RLock()
-		defer ls.LivepeerNode.LiveMu.RUnlock()
+		// NB: LiveMu is a global lock, avoid holding it
+		// during blocking network actions ... can't defer
 		p, ok := ls.LivepeerNode.LivePipelines[stream]
+		ls.LivepeerNode.LiveMu.RUnlock()
 		if !ok {
 			// Stream not found
 			http.Error(w, "Stream not found", http.StatusNotFound)
 			return
 		}
-		defer r.Body.Close()
-		params, err := io.ReadAll(r.Body)
+		reader := http.MaxBytesReader(w, r.Body, 1*1024*104) // 1 MB
+		defer reader.Close()
+		data, err := io.ReadAll(r.Body)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
 
-		clog.V(6).Infof(ctx, "Sending Live Video Update Control API stream=%s, params=%s", stream, string(params))
-		if err := p.ControlPub.Write(strings.NewReader(string(params))); err != nil {
+		params := string(data)
+		ls.LivepeerNode.LiveMu.Lock()
+		p.Params = params
+		controlPub := p.ControlPub
+		ls.LivepeerNode.LiveMu.Unlock()
+
+		if controlPub == nil {
+			clog.Info(ctx, "No orchestrator available, caching params", "stream", stream, "params", params)
+			return
+		}
+
+		clog.V(6).Infof(ctx, "Sending Live Video Update Control API stream=%s, params=%s", stream, params)
+		if err := controlPub.Write(strings.NewReader(params)); err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
@@ -1069,6 +1089,8 @@ func (ls *LivepeerServer) CreateWhip(server *media.WHIPServer) http.Handler {
 				},
 			}
 
+			registerControl(ctx, params)
+
 			req := worker.GenLiveVideoToVideoJSONRequestBody{
 				ModelId:          &pipeline,
 				Params:           &pipelineParams,