diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 5ecfc8379c..56813d0d07 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -82,6 +82,7 @@ jobs: with: build-args: | BUILD_TAGS=${{ steps.build-tag.outputs.build-tags }} + GIT_REVISION=${{ github.event.pull_request.head.sha || github.sha }} context: . provenance: mode=max sbom: true @@ -92,8 +93,10 @@ jobs: tags: ${{ steps.meta.outputs.tags }} annotations: ${{ steps.meta.outputs.annotations }} labels: ${{ steps.meta.outputs.labels }} - cache-from: type=registry,ref=livepeerci/build:cache - cache-to: type=registry,ref=livepeerci/build:cache,mode=max + # PR builds must not reuse the shared registry cache — it can serve a + # stale livepeer binary while metadata tags the PR head SHA (see #3947). + cache-from: ${{ github.event_name != 'pull_request' && 'type=registry,ref=livepeerci/build:cache' || '' }} + cache-to: ${{ github.event_name != 'pull_request' && 'type=registry,ref=livepeerci/build:cache,mode=max' || '' }} builder: name: go-livepeer builder docker image generation diff --git a/ai/runner/live_runner.go b/ai/runner/live_runner.go new file mode 100644 index 0000000000..e0b80df8d0 --- /dev/null +++ b/ai/runner/live_runner.go @@ -0,0 +1,1397 @@ +package runner + +import ( + "crypto/hmac" + "crypto/rand" + "crypto/sha256" + "encoding/base32" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math/big" + "net/http" + "net/url" + "regexp" + "sort" + "strings" + "sync" + "time" + + "github.com/jaypipes/ghw" + "github.com/livepeer/go-livepeer/common" + "github.com/livepeer/go-livepeer/core" + "github.com/livepeer/go-livepeer/trickle" +) + +const ( + defaultLiveRunnerHeartbeatInterval = 5 * time.Second + defaultLiveRunnerHeartbeatTTL = 20 * time.Second + defaultLiveRunnerHealthInterval = 5 * time.Second + defaultLiveRunnerHealthStatusCode = http.StatusOK + defaultLiveRunnerHealthTimeout = 3 * time.Second + liveRunnerO2RKeepaliveMessage = `{"keep":"alive"}` + liveRunnerIDRandomBytes = 5 + liveRunnerSeedBytes = 32 +) + +var liveRunnerO2RKeepaliveInterval = 10 * time.Second + +const ( + LiveRunnerModePersistent = "persistent" + LiveRunnerModeSingleShot = "single-shot" + liveRunnerModeSingleShotAlias = "single_shot" +) + +const ( + LiveRunnerRoutingRunnerID = "runner-id" + LiveRunnerRoutingLabel = "label" +) + +const ( + // o2r gives the orchestrator an inbound signaling path to the runner without + // requiring a new runner port or quietly extending the runner app API. + // Runner-originated calls still use normal HTTP endpoints for heartbeats, + // trickle channel creation, and related control-plane requests. + LiveRunnerTrickleOrchestratorToRunner = "o2r" +) + +var liveRunnerEncoding = base32.StdEncoding.WithPadding(base32.NoPadding) + +type RunnerError struct { + StatusCode int + Message string +} + +func (e *RunnerError) Error() string { + if e == nil { + return "runner error" + } + return e.Message +} + +type LiveRunnerGPU struct { + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + VRAMMB int `json:"vram_mb,omitempty"` +} + +type LiveRunnerPriceInfo struct { + PricePerUnit int64 `json:"price_per_unit"` + PixelsPerUnit int64 `json:"pixels_per_unit"` + Unit string `json:"unit,omitempty"` +} + +type LiveRunnerHeartbeatRequest struct { + RunnerID string `json:"runner_id,omitempty"` + Label string `json:"label,omitempty"` + RunnerURL string `json:"runner_url"` + Version string `json:"version,omitempty"` + Status string `json:"status,omitempty"` + Mode string `json:"mode,omitempty"` + GPU *LiveRunnerGPU `json:"gpu,omitempty"` + App string `json:"app"` + Capacity int `json:"capacity"` + PriceInfo LiveRunnerPriceInfo `json:"price_info"` + SessionIDs []string `json:"session_ids,omitempty"` +} + +type StaticLiveRunnerConfig struct { + Runners []StaticLiveRunnerConfigEntry `json:"runners"` +} + +type StaticLiveRunnerConfigEntry struct { + Label string `json:"label,omitempty"` + Route string `json:"routing,omitempty"` + RunnerURL string `json:"runner_url"` + Version string `json:"version,omitempty"` + Mode string `json:"mode,omitempty"` + GPU *LiveRunnerGPU `json:"gpu,omitempty"` + App string `json:"app"` + Capacity int `json:"capacity"` + PriceInfo LiveRunnerPriceInfo `json:"price_info"` + HealthURL string `json:"health_url"` + HealthyStatusCode int `json:"healthy_status_code,omitempty"` +} + +type StaticLiveRunnerRegistration struct { + RunnerID string `json:"runner_id"` + Label string `json:"label,omitempty"` + App string `json:"app"` + RunnerURL string `json:"runner_url"` + Mode string `json:"mode,omitempty"` + HealthURL string `json:"health_url"` + HealthyStatusCode int `json:"healthy_status_code"` + Healthy bool `json:"healthy"` +} + +type StaticLiveRunnerRegistrationResponse struct { + Runners []StaticLiveRunnerRegistration `json:"runners"` +} + +type LiveRunnerHeartbeatResponse struct { + RunnerID string `json:"runner_id"` + Orchestrator string `json:"orchestrator,omitempty"` + HeartbeatInterval string `json:"heartbeat_interval"` + HeartbeatTTL string `json:"heartbeat_ttl"` + HeartbeatSecret string `json:"heartbeat_secret,omitempty"` + O2R *LiveRunnerTrickleChannel `json:"o2r,omitempty"` + SessionIDs []string `json:"session_ids"` +} + +type LiveRunnerDiscoveryRunner struct { + URL string `json:"url"` + GPU *LiveRunnerGPU `json:"gpu,omitempty"` + App string `json:"app"` + Version string `json:"version,omitempty"` + Mode string `json:"mode,omitempty"` + Capacity int `json:"capacity"` + CapacityUsed int `json:"capacity_used"` + CapacityAvailable int `json:"capacity_available"` + PriceInfo *LiveRunnerPriceInfo `json:"price_info,omitempty"` +} + +type LiveRunnerTrickleChannel struct { + Name string `json:"name"` + ChannelName string `json:"channel_name"` + URL string `json:"url"` + InternalURL string `json:"internal_url,omitempty"` + MimeType string `json:"mime_type"` +} + +type liveRunner struct { + seed []byte + heartbeatSecret string + static bool + offchain bool + o2r *liveRunnerTrickleChannel + + // all fields below are mutable and need to be protected by mu + mu sync.Mutex + LiveRunnerHeartbeatRequest + LastHeartbeat time.Time + removed bool + healthURL string + healthStatus int + route string + sessions map[string]*liveRunnerSession + priceSource LiveRunnerPriceInfo + converter *core.AutoConvertedPrice +} + +type liveRunnerSession struct { + // Protected by the parent liveRunner.mu. + createdAt time.Time + channels map[string]*liveRunnerTrickleChannel +} + +type liveRunnerTrickleChannel struct { + channel LiveRunnerTrickleChannel + publisher *trickle.TrickleLocalPublisher + done chan struct{} + messages chan string + closeOnce sync.Once +} + +type LiveRunnerRegistry struct { + host RunnerHost + offchain bool + heartbeatInterval time.Duration + heartbeatTTL time.Duration + healthClient *http.Client + healthInterval time.Duration + stopRegistry chan struct{} + stopRegistryOnce sync.Once + + staticRegistrationMu sync.Mutex // serializes static runner config upserts + + // mu protects runner registration, route lookup, and trickle server routing. + mu sync.Mutex + runners map[string]*liveRunner + trickleSrv *trickle.Server + publicTrickleBaseURL string + internalTrickleBaseURL string +} + +type RunnerHost interface { + ServiceURI() *url.URL + LiveRunnerURI() *url.URL + RegistrationSecret() string +} + +type LiveRunnerRegistryConfig struct { + Host RunnerHost + Onchain bool +} + +var liveRunnerChannelNamePattern = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) + +func NewLiveRunnerRegistry(config LiveRunnerRegistryConfig) *LiveRunnerRegistry { + r := &LiveRunnerRegistry{ + host: config.Host, + offchain: !config.Onchain, + runners: make(map[string]*liveRunner), + heartbeatInterval: defaultLiveRunnerHeartbeatInterval, + heartbeatTTL: defaultLiveRunnerHeartbeatTTL, + healthClient: &http.Client{Timeout: defaultLiveRunnerHealthTimeout}, + healthInterval: defaultLiveRunnerHealthInterval, + stopRegistry: make(chan struct{}), + } + go r.healthLoop() + go r.expiryLoop() + return r +} + +func (r *LiveRunnerRegistry) SetTrickleServer(srv *trickle.Server, publicBaseURL, internalBaseURL string) { + r.mu.Lock() + defer r.mu.Unlock() + r.trickleSrv = srv + r.publicTrickleBaseURL = publicBaseURL + r.internalTrickleBaseURL = internalBaseURL +} + +func (r *LiveRunnerRegistry) Stop() { + if r == nil { + return + } + r.stopRegistryOnce.Do(func() { + close(r.stopRegistry) + }) + r.mu.Lock() + runners := make([]*liveRunner, 0, len(r.runners)) + for _, runner := range r.runners { + runners = append(runners, runner) + } + r.mu.Unlock() + for _, runner := range runners { + runner.mu.Lock() + runner.closeChannelsLocked() + runner.stopPriceConverterLocked() + runner.mu.Unlock() + } +} + +func (r *LiveRunnerRegistry) Heartbeat(req LiveRunnerHeartbeatRequest, auth string) (*LiveRunnerHeartbeatResponse, error) { + req.RunnerID = strings.TrimSpace(req.RunnerID) + if req.RunnerID == "" { + req.RunnerID = "runner_" + randomStr(liveRunnerIDRandomBytes) + } + + req, err := r.normalizeHeartbeat(req.RunnerID, req) + if err != nil { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: err.Error()} + } + + // For now, only return the heartbeat secret on first registration. + // Return another one later on if the secret needs to be rotated. + var responseHeartbeatSecret string + var runner *liveRunner + + r.mu.Lock() + runner = r.runners[req.RunnerID] + heartbeatInterval := r.heartbeatInterval + heartbeatTTL := r.heartbeatTTL + orchestrator := r.host.LiveRunnerURI().String() + if runner == nil { + trickleSrv := r.trickleSrv + trickleBaseURL := r.internalTrickleBaseURL + // registering a new runner + registrationSecret := r.host.RegistrationSecret() + if registrationSecret == "" { + r.mu.Unlock() + return nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "dynamic live runner registration is disabled"} + } + if !validSecret(auth, registrationSecret) { + r.mu.Unlock() + return nil, &RunnerError{StatusCode: http.StatusUnauthorized, Message: "invalid registration authorization"} + } + seed := newRunnerSeed() + heartbeatSecret := deriveSecret(seed, "heartbeat") + runner = &liveRunner{ + seed: seed, + heartbeatSecret: heartbeatSecret, + offchain: r.offchain, + route: req.RunnerID, // No label-based routing for now + sessions: make(map[string]*liveRunnerSession), + } + if err := runner.setRunnerTrickleChannels(trickleSrv, trickleBaseURL, req.RunnerID); err != nil { + runner.closeChannelsLocked() + r.mu.Unlock() + return nil, err + } + r.runners[req.RunnerID] = runner + responseHeartbeatSecret = heartbeatSecret + runner.mu.Lock() + r.mu.Unlock() + + defer runner.mu.Unlock() + runner.LiveRunnerHeartbeatRequest = req + runner.LastHeartbeat = time.Now() + runner.route = runner.RunnerID // not doing label based routing here for now + runner.updatePriceConverterLocked() + + // SessionIDs in the response come from the orchestrator's authoritative + // active sessions and may differ from session_ids reported in the request. + return &LiveRunnerHeartbeatResponse{ + RunnerID: runner.RunnerID, + Orchestrator: orchestrator, + HeartbeatInterval: heartbeatInterval.String(), + HeartbeatTTL: heartbeatTTL.String(), + HeartbeatSecret: responseHeartbeatSecret, + O2R: &runner.o2r.channel, + SessionIDs: runner.sessionIDsLocked(), + }, nil + } else { + r.mu.Unlock() + } + + runner.mu.Lock() + defer runner.mu.Unlock() + if runner.removed || liveRunnerExpiredLocked(runner, time.Now(), heartbeatTTL) { + return nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if !validSecret(auth, runner.heartbeatSecret) { + return nil, &RunnerError{StatusCode: http.StatusUnauthorized, Message: "invalid heartbeat authorization"} + } + + if runner.static { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: "static runner cannot heartbeat"} + } + runner.LiveRunnerHeartbeatRequest = req + runner.LastHeartbeat = time.Now() + runner.route = runner.RunnerID // not doing label based routing here for now + runner.updatePriceConverterLocked() + + // SessionIDs in the response come from the orchestrator's authoritative + // active sessions and may differ from session_ids reported in the request. + return &LiveRunnerHeartbeatResponse{ + RunnerID: runner.RunnerID, + Orchestrator: orchestrator, + HeartbeatInterval: heartbeatInterval.String(), + HeartbeatTTL: heartbeatTTL.String(), + HeartbeatSecret: responseHeartbeatSecret, + SessionIDs: runner.sessionIDsLocked(), + }, nil +} + +func ParseStaticLiveRunnerConfig(data []byte) (StaticLiveRunnerConfig, error) { + var cfg StaticLiveRunnerConfig + if err := json.Unmarshal(data, &cfg); err != nil { + return StaticLiveRunnerConfig{}, err + } + return normalizeStaticLiveRunnerConfig(cfg) +} + +func (r *LiveRunnerRegistry) RegisterStaticRunnersJSON(data []byte) (*StaticLiveRunnerRegistrationResponse, error) { + cfg, err := ParseStaticLiveRunnerConfig(data) + if err != nil { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: err.Error()} + } + return r.RegisterStaticRunners(cfg) +} + +func normalizeStaticLiveRunnerConfig(cfg StaticLiveRunnerConfig) (StaticLiveRunnerConfig, error) { + for i := range cfg.Runners { + cfg.Runners[i].Route = strings.TrimSpace(cfg.Runners[i].Route) + if cfg.Runners[i].Route == "" { + cfg.Runners[i].Route = LiveRunnerRoutingRunnerID + } + if cfg.Runners[i].Route != LiveRunnerRoutingRunnerID && cfg.Runners[i].Route != LiveRunnerRoutingLabel { + return StaticLiveRunnerConfig{}, fmt.Errorf("runners[%d]: routing must be %q or %q", i, LiveRunnerRoutingRunnerID, LiveRunnerRoutingLabel) + } + if cfg.Runners[i].HealthyStatusCode == 0 { + cfg.Runners[i].HealthyStatusCode = defaultLiveRunnerHealthStatusCode + } + healthURL, err := staticRunnerHealthURL(cfg.Runners[i].RunnerURL, cfg.Runners[i].HealthURL) + if err != nil { + return StaticLiveRunnerConfig{}, fmt.Errorf("runners[%d]: %v", i, err) + } + cfg.Runners[i].HealthURL = healthURL + } + return cfg, nil +} + +func staticRunnerHealthURL(runnerURL, healthURL string) (string, error) { + u, err := url.Parse(healthURL) + if err != nil { + return "", fmt.Errorf("invalid health_url") + } + if u.IsAbs() { + return healthURL, nil + } + if u.Path == "" || u.Path[0] != '/' || u.Host != "" { + return healthURL, nil + } + base, err := url.ParseRequestURI(runnerURL) + if err != nil || base.Host == "" || (base.Scheme != "http" && base.Scheme != "https") { + return "", fmt.Errorf("invalid runner_url") + } + return strings.TrimRight(runnerURL, "/") + healthURL, nil +} + +func staticRunnerRoute(route, runnerID, label string) string { + if route == LiveRunnerRoutingLabel { + return label + } + return runnerID +} + +func (g *LiveRunnerGPU) UnmarshalJSON(data []byte) error { + var deviceIndex int + if err := json.Unmarshal(data, &deviceIndex); err == nil { + *g = liveRunnerGPUForIndex(deviceIndex) + return nil + } + type gpuAlias LiveRunnerGPU + var parsed gpuAlias + if err := json.Unmarshal(data, &parsed); err != nil { + return err + } + *g = LiveRunnerGPU(parsed) + return nil +} + +func liveRunnerGPUForIndex(deviceIndex int) LiveRunnerGPU { + gpuInfo := LiveRunnerGPU{ID: fmt.Sprintf("%d", deviceIndex)} + if deviceIndex < 0 { + return gpuInfo + } + info, err := ghw.GPU() + if err != nil || info == nil || deviceIndex >= len(info.GraphicsCards) { + return gpuInfo + } + card := info.GraphicsCards[deviceIndex] + if card == nil { + return gpuInfo + } + if card.DeviceInfo != nil { + gpuInfo.Name = strings.TrimSpace(card.DeviceInfo.Product.Name) + } + return gpuInfo +} + +func (r *LiveRunnerRegistry) RegisterStaticRunners(cfg StaticLiveRunnerConfig) (*StaticLiveRunnerRegistrationResponse, error) { + cfg, err := normalizeStaticLiveRunnerConfig(cfg) + if err != nil { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: err.Error()} + } + requests := make([]LiveRunnerHeartbeatRequest, 0, len(cfg.Runners)) + health := make([]bool, 0, len(cfg.Runners)) + seenLabels := make(map[string]struct{}, len(cfg.Runners)) + for i, entry := range cfg.Runners { + req, err := r.buildStaticRunner(entry) + if err != nil { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: fmt.Sprintf("runners[%d]: %v", i, err)} + } + if _, duplicate := seenLabels[entry.Label]; duplicate { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: fmt.Sprintf("runners[%d]: duplicate label %q", i, entry.Label)} + } + if entry.Route == LiveRunnerRoutingLabel && strings.Contains(entry.Label, "/") { + return nil, &RunnerError{StatusCode: http.StatusBadRequest, Message: fmt.Sprintf("runners[%d]: label cannot contain / when routing is %q", i, LiveRunnerRoutingLabel)} + } + seenLabels[entry.Label] = struct{}{} + requests = append(requests, req) + health = append(health, r.staticRunnerHealthy(entry.HealthURL, entry.HealthyStatusCode)) + } + + r.staticRegistrationMu.Lock() + defer r.staticRegistrationMu.Unlock() + + r.mu.Lock() + snapshot := make([]*liveRunner, 0, len(r.runners)) + for _, runner := range r.runners { + snapshot = append(snapshot, runner) + } + r.mu.Unlock() + + existingStatic := make(map[string]*liveRunner) + for _, runner := range snapshot { + runner.mu.Lock() + if !runner.removed && runner.static && runner.Label != "" { + existingStatic[runner.Label] = runner + } + runner.mu.Unlock() + } + + staticRunnersByLabel := make(map[string]*liveRunner, len(cfg.Runners)) + for _, entry := range cfg.Runners { + staticRunner := existingStatic[entry.Label] + if staticRunner == nil { + staticRunner = &liveRunner{ + seed: newRunnerSeed(), + static: true, + offchain: r.offchain, + sessions: make(map[string]*liveRunnerSession), + } + staticRunner.RunnerID = "runner_" + randomStr(liveRunnerIDRandomBytes) + } + staticRunnersByLabel[entry.Label] = staticRunner + } + + registrations := make([]StaticLiveRunnerRegistration, 0, len(cfg.Runners)) + for i, entry := range cfg.Runners { + staticRunner := staticRunnersByLabel[entry.Label] + staticRunner.mu.Lock() + r.mu.Lock() + if r.runners[staticRunner.RunnerID] != nil && r.runners[staticRunner.RunnerID] != staticRunner { + for r.runners[staticRunner.RunnerID] != nil { + // collision lol + staticRunner.RunnerID = "runner_" + randomStr(liveRunnerIDRandomBytes) + } + } + req := requests[i] + req.RunnerID = staticRunner.RunnerID + if health[i] { + req.Status = "ready" + } else { + req.Status = "unavailable" + for sessionID := range staticRunner.sessions { + staticRunner.releaseSessionLocked(sessionID) + } + } + staticRunner.removed = false + staticRunner.LiveRunnerHeartbeatRequest = req + staticRunner.LastHeartbeat = time.Now() + staticRunner.healthURL = entry.HealthURL + staticRunner.healthStatus = entry.HealthyStatusCode + staticRunner.route = staticRunnerRoute(entry.Route, staticRunner.RunnerID, req.Label) + for route, runner := range r.runners { + if runner == staticRunner && route != staticRunner.route { + // route changed so remove old entry and re-add + delete(r.runners, route) + } + } + r.runners[staticRunner.route] = staticRunner + staticRunner.updatePriceConverterLocked() + r.mu.Unlock() + registrations = append(registrations, StaticLiveRunnerRegistration{ + RunnerID: staticRunner.RunnerID, + Label: req.Label, + App: req.App, + RunnerURL: req.RunnerURL, + Mode: req.Mode, + HealthURL: entry.HealthURL, + HealthyStatusCode: entry.HealthyStatusCode, + Healthy: health[i], + }) + staticRunner.mu.Unlock() + } + return &StaticLiveRunnerRegistrationResponse{Runners: registrations}, nil +} + +func (r *LiveRunnerRegistry) buildStaticRunner(entry StaticLiveRunnerConfigEntry) (LiveRunnerHeartbeatRequest, error) { + if entry.Label == "" || entry.Label != strings.TrimSpace(entry.Label) { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("label is required and must be trimmed") + } + if entry.HealthURL == "" { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("health_url is required") + } + u, err := url.ParseRequestURI(entry.HealthURL) + if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("invalid health_url") + } + if entry.HealthyStatusCode < 100 || entry.HealthyStatusCode > 599 { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("healthy_status_code must be a valid HTTP status code") + } + req := LiveRunnerHeartbeatRequest{ + Label: entry.Label, + RunnerURL: entry.RunnerURL, + Version: entry.Version, + Mode: entry.Mode, + GPU: entry.GPU, + App: entry.App, + Capacity: entry.Capacity, + PriceInfo: entry.PriceInfo, + } + return r.normalizeHeartbeat("static", req) +} + +func (r *LiveRunnerRegistry) normalizeHeartbeat(runnerID string, req LiveRunnerHeartbeatRequest) (LiveRunnerHeartbeatRequest, error) { + u, err := url.ParseRequestURI(req.RunnerURL) + if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("invalid runner_url") + } + + if req.App == "" || req.App != strings.TrimSpace(req.App) { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("app must be trimmed") + } + if req.Status != strings.ToLower(strings.TrimSpace(req.Status)) { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("status must be lowercase and trimmed") + } + if req.Mode == "" { + req.Mode = LiveRunnerModePersistent + } + if req.Mode == liveRunnerModeSingleShotAlias { + req.Mode = LiveRunnerModeSingleShot + } + if req.Mode != LiveRunnerModePersistent && req.Mode != LiveRunnerModeSingleShot { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("mode must be %q or %q", LiveRunnerModePersistent, LiveRunnerModeSingleShot) + } + if req.Capacity < 0 { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("capacity must be >= 0") + } + if req.Capacity == 0 { + req.Capacity = 1 + } + req.RunnerID = runnerID + req.GPU = cloneLiveRunnerGPU(req.GPU) + + if r.offchain { + return req, nil + } + + if req.PriceInfo.PixelsPerUnit <= 0 || req.PriceInfo.PricePerUnit <= 0 { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("runner must include positive price_info") + } + unit := strings.ToUpper(strings.TrimSpace(req.PriceInfo.Unit)) + if unit != "USD" && unit != "WEI" { + return LiveRunnerHeartbeatRequest{}, fmt.Errorf("price_info.unit must be USD or WEI") + } + req.PriceInfo.Unit = unit + + return req, nil +} + +func (r *LiveRunnerRegistry) Unregister(runnerID, auth string) error { + r.mu.Lock() + runner := r.runners[runnerID] + heartbeatTTL := r.heartbeatTTL + if runner == nil { + r.mu.Unlock() + return &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + r.mu.Unlock() + + runner.mu.Lock() + defer runner.mu.Unlock() + if runner.removed || liveRunnerExpiredLocked(runner, time.Now(), heartbeatTTL) { + return &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if runner.static { + return &RunnerError{StatusCode: http.StatusBadRequest, Message: "static runner cannot unregister with heartbeat credentials"} + } + if !validSecret(auth, runner.heartbeatSecret) { + return &RunnerError{StatusCode: http.StatusUnauthorized, Message: "invalid authorization"} + } + r.mu.Lock() + if r.runners[runnerID] == runner && !runner.removed { + r.removeRunnerWithRunnerLocked(runnerID, runner) + } + r.mu.Unlock() + return nil +} + +func (r *LiveRunnerRegistry) ReserveSession(runnerID string, optSessionID ...string) (string, string, error) { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return "", "", err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return "", "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if len(runner.sessions) >= runner.Capacity { + return "", "", &RunnerError{StatusCode: http.StatusServiceUnavailable, Message: "no capacity available for runner"} + } + id := "" + if len(optSessionID) > 0 { + // use supplied id if given - usually ticket param's auth token id / manifest id + id = optSessionID[0] + } + if id != "" { + if _, exists := runner.sessions[id]; exists { + return "", "", &RunnerError{StatusCode: http.StatusConflict, Message: "session already exists"} + } + } else { + id = "session_" + randomStr(liveRunnerIDRandomBytes) + for { + if _, exists := runner.sessions[id]; !exists { + break + } + id = "session_" + randomStr(liveRunnerIDRandomBytes) + } + } + runner.sessions[id] = &liveRunnerSession{ + createdAt: time.Now(), + channels: make(map[string]*liveRunnerTrickleChannel), + } + runner.sendSessionEvent("reserved", id) + return id, runner.RunnerURL, nil +} + +func (r *LiveRunnerRegistry) PaymentInfo(runnerID string) (*LiveRunnerPriceInfo, error) { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return nil, err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if runner.offchain { + return nil, nil + } + priceInfo := runner.PriceInfo + if runner.converter != nil { + converted, err := convertedPriceInfo(runner.converter) + if err != nil { + return nil, err + } + priceInfo = converted + } + if priceInfo == (LiveRunnerPriceInfo{}) || priceInfo.PricePerUnit <= 0 || priceInfo.PixelsPerUnit <= 0 { + return nil, nil + } + return &priceInfo, nil +} + +func (r *LiveRunnerRegistry) ReleaseSession(runnerID, sessionID string) error { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return err + } + defer unlock() + runner.releaseSessionLocked(sessionID) + return nil +} + +func (r *LiveRunnerRegistry) RunnerMode(runnerID string) (string, error) { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return "", err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if runner.Mode == "" { + return LiveRunnerModePersistent, nil + } + return runner.Mode, nil +} + +func (r *LiveRunnerRegistry) RunnerEndpointForSession(runnerID, sessionID string) (string, error) { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return "", err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if _, exists := runner.sessions[sessionID]; !exists { + return "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner session not found"} + } + return runner.RunnerURL, nil +} + +func (r *LiveRunnerRegistry) SessionTokenForSession(runnerID, sessionID string) (string, error) { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return "", err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if _, exists := runner.sessions[sessionID]; !exists { + return "", &RunnerError{StatusCode: http.StatusNotFound, Message: "runner session not found"} + } + return runner.sessionToken(sessionID), nil +} + +func (r *LiveRunnerRegistry) ValidSessionToken(runnerID, sessionID, token string) error { + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return err + } + defer unlock() + if !isReadyStatus(runner.Status) { + return &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + if _, exists := runner.sessions[sessionID]; !exists { + return &RunnerError{StatusCode: http.StatusNotFound, Message: "runner session not found"} + } + if !validSecret(token, runner.sessionToken(sessionID)) { + return &RunnerError{StatusCode: http.StatusUnauthorized, Message: "invalid authorization"} + } + return nil +} + +func (r *LiveRunnerRegistry) CreateTrickleChannel(runnerID, sessionID, name, mimeType string) (LiveRunnerTrickleChannel, error) { + name, channelName, err := normalizeTrickleChannelName(sessionID, name) + if err != nil { + return LiveRunnerTrickleChannel{}, err + } + if strings.TrimSpace(mimeType) == "" { + mimeType = "application/octet-stream" + } + + r.mu.Lock() + trickleSrv := r.trickleSrv + publicTrickleBaseURL := r.publicTrickleBaseURL + internalTrickleBaseURL := r.internalTrickleBaseURL + r.mu.Unlock() + + if trickleSrv == nil { + return LiveRunnerTrickleChannel{}, fmt.Errorf("trickle server is required") + } + channelURL, err := url.JoinPath(publicTrickleBaseURL, channelName) + if err != nil { + return LiveRunnerTrickleChannel{}, err + } + internalChannelURL, err := url.JoinPath(internalTrickleBaseURL, channelName) + if err != nil { + return LiveRunnerTrickleChannel{}, err + } + + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return LiveRunnerTrickleChannel{}, err + } + defer unlock() + session, err := runner.liveRunnerSessionLocked(sessionID) + if err != nil { + return LiveRunnerTrickleChannel{}, err + } + if channel := session.channels[name]; channel != nil { + return channel.channel, nil + } + + channel := newLiveRunnerTrickleChannel(trickleSrv, name, channelName, channelURL, internalChannelURL, mimeType) + session.channels[name] = channel + return channel.channel, nil +} + +func (runner *liveRunner) setRunnerTrickleChannels(trickleSrv *trickle.Server, trickleBaseURL, runnerID string) error { + o2r, err := runner.createBootstrapTrickleChannel(trickleSrv, trickleBaseURL, runnerID, LiveRunnerTrickleOrchestratorToRunner) + if err != nil { + return err + } + runner.o2r = o2r + go writeO2RLoop(runner.o2r) + return nil +} + +func (runner *liveRunner) createBootstrapTrickleChannel(trickleSrv *trickle.Server, trickleBaseURL, runnerID, name string) (*liveRunnerTrickleChannel, error) { + channelName := runnerID + "-" + deriveSecret(runner.seed, "trickle:"+name)[:16] + "-" + name + channelURL, err := url.JoinPath(trickleBaseURL, channelName) + if err != nil { + return nil, err + } + channel := newLiveRunnerTrickleChannel(trickleSrv, name, channelName, channelURL, "", "application/json") + return channel, nil +} + +func newLiveRunnerTrickleChannel(trickleSrv *trickle.Server, name, channelName, channelURL, internalChannelURL, mimeType string) *liveRunnerTrickleChannel { + publisher := trickle.NewLocalPublisher(trickleSrv, channelName, mimeType) + publisher.CreateChannel() + return &liveRunnerTrickleChannel{ + channel: LiveRunnerTrickleChannel{ + Name: name, + ChannelName: channelName, + URL: channelURL, + InternalURL: internalChannelURL, + MimeType: mimeType, + }, + publisher: publisher, + done: make(chan struct{}), + messages: make(chan string, 32), + } +} + +func (r *LiveRunnerRegistry) DeleteTrickleChannel(runnerID, sessionID, name string) error { + name, _, err := normalizeTrickleChannelName(sessionID, name) + if err != nil { + return err + } + + runner, unlock, err := r.lockLiveRunner(runnerID) + if err != nil { + return err + } + defer unlock() + session, err := runner.liveRunnerSessionLocked(sessionID) + if err != nil { + return err + } + channel := session.channels[name] + if channel == nil { + return &RunnerError{StatusCode: http.StatusNotFound, Message: "trickle channel not found"} + } + if err := channel.close(); err != nil { + return err + } + delete(session.channels, name) + return nil +} + +func (r *LiveRunnerRegistry) Runners() []LiveRunnerDiscoveryRunner { + r.mu.Lock() + heartbeatTTL := r.heartbeatTTL + snapshot := make([]*liveRunner, 0, len(r.runners)) + for _, runner := range r.runners { + snapshot = append(snapshot, runner) + } + r.mu.Unlock() + + runners := make([]LiveRunnerDiscoveryRunner, 0, len(snapshot)) + for _, runner := range snapshot { + runner.mu.Lock() + if runner.removed || liveRunnerExpiredLocked(runner, time.Now(), heartbeatTTL) || !isReadyStatus(runner.Status) { + runner.mu.Unlock() + continue + } + discoveryRunner := runner.discoveryRunner() + discoveryRunner.URL = r.discoveryRunnerURL(runner) + runner.mu.Unlock() + runners = append(runners, discoveryRunner) + } + return runners +} + +func (r *LiveRunnerRegistry) discoveryRunnerURL(runner *liveRunner) string { + // Under no circumstances leak the runner URL through discovery. + // Discovery must only expose orchestrator proxy URLs; the runner URL is an internal endpoint. + // discoveryRunnerURL requires runner.mu. + if runner.Mode == LiveRunnerModeSingleShot { + return r.host.ServiceURI().JoinPath("apps", runner.route, "app").String() + } + return r.host.ServiceURI().JoinPath("apps", runner.route, "session").String() +} + +func (r *LiveRunnerRegistry) lockLiveRunner(runnerID string) (*liveRunner, func(), error) { + r.mu.Lock() + runner := r.runners[runnerID] + if runner == nil { + r.mu.Unlock() + return nil, nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + heartbeatTTL := r.heartbeatTTL + r.mu.Unlock() + + // Do not wait on runner.mu while holding r.mu. A copied runner pointer stays + // valid after map deletion; removal is observed through runner.removed once + // runner.mu is acquired. + runner.mu.Lock() + if runner.removed || liveRunnerExpiredLocked(runner, time.Now(), heartbeatTTL) { + runner.mu.Unlock() + return nil, nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + return runner, runner.mu.Unlock, nil +} + +func liveRunnerExpiredLocked(runner *liveRunner, now time.Time, heartbeatTTL time.Duration) bool { + // liveRunnerExpiredLocked requires runner.mu. + return !runner.static && now.Sub(runner.LastHeartbeat) > heartbeatTTL +} + +func (r *LiveRunnerRegistry) expiryLoop() { + ticker := time.NewTicker(r.heartbeatInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + r.removeExpiredRunners(time.Now()) + case <-r.stopRegistry: + return + } + } +} + +func (r *LiveRunnerRegistry) removeExpiredRunners(now time.Time) { + type candidate struct { + runnerID string + runner *liveRunner + } + r.mu.Lock() + heartbeatTTL := r.heartbeatTTL + candidates := make([]candidate, 0, len(r.runners)) + for runnerID, runner := range r.runners { + candidates = append(candidates, candidate{runnerID: runnerID, runner: runner}) + } + r.mu.Unlock() + + for _, c := range candidates { + c.runner.mu.Lock() + expired := !c.runner.removed && liveRunnerExpiredLocked(c.runner, now, heartbeatTTL) + if expired { + r.mu.Lock() + if r.runners[c.runnerID] == c.runner && !c.runner.removed && liveRunnerExpiredLocked(c.runner, now, heartbeatTTL) { + r.removeRunnerWithRunnerLocked(c.runnerID, c.runner) + } + r.mu.Unlock() + } + c.runner.mu.Unlock() + } +} + +func (r *LiveRunnerRegistry) healthLoop() { + ticker := time.NewTicker(r.healthInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + r.checkStaticRunnerHealth() + case <-r.stopRegistry: + return + } + } +} + +func (r *LiveRunnerRegistry) checkStaticRunnerHealth() { + type check struct { + runner *liveRunner + healthURL string + healthStatus int + } + r.mu.Lock() + checks := []check{} + for _, runner := range r.runners { + checks = append(checks, check{runner: runner}) + } + r.mu.Unlock() + + for i, c := range checks { + runner := c.runner + runner.mu.Lock() + if !runner.removed && runner.static { + checks[i].healthURL = runner.healthURL + checks[i].healthStatus = runner.healthStatus + } else { + checks[i].runner = nil + } + runner.mu.Unlock() + } + + for _, c := range checks { + if c.runner == nil { + continue + } + healthy := r.staticRunnerHealthy(c.healthURL, c.healthStatus) + c.runner.mu.Lock() + if !c.runner.removed && c.runner.static { + if healthy { + c.runner.Status = "ready" + } else { + c.runner.Status = "unavailable" + for sessionID := range c.runner.sessions { + c.runner.releaseSessionLocked(sessionID) + } + } + } + c.runner.mu.Unlock() + } +} + +func (r *LiveRunnerRegistry) staticRunnerHealthy(healthURL string, healthyStatus int) bool { + resp, err := r.healthClient.Get(healthURL) + if err != nil { + return false + } + defer resp.Body.Close() + return resp.StatusCode == healthyStatus +} + +func (r *LiveRunnerRegistry) removeRunnerWithRunnerLocked(runnerID string, runner *liveRunner) { + // removeRunnerWithRunnerLocked requires both r.mu and runner.mu. Callers + // should already hold runner.mu before taking r.mu so the registry lock is + // not held while waiting on a busy runner. + // Set removed before deleting the map entry so any goroutine that already + // copied the runner pointer can reject it after acquiring runner.mu. + runner.removed = true + runner.closeChannelsLocked() + for sessionID := range runner.sessions { + runner.releaseSessionLocked(sessionID) + } + runner.stopPriceConverterLocked() + if r.runners[runnerID] == runner { + delete(r.runners, runnerID) + } +} + +func (runner *liveRunner) releaseSessionLocked(sessionID string) { + // releaseSessionLocked requires runner.mu. + if session := runner.sessions[sessionID]; session != nil { + runner.sendSessionEvent("released", sessionID) + session.closeChannelsLocked() + } + delete(runner.sessions, sessionID) +} + +func (runner *liveRunner) sendSessionEvent(event, sessionID string) { + o2r := runner.o2r + if o2r == nil { + return + } + if err := o2r.queueJSON(struct { + Event string `json:"event"` + Session string `json:"session"` + Timestamp time.Time `json:"timestamp"` + }{ + Event: event, + Session: sessionID, + Timestamp: time.Now().UTC(), + }); err != nil { + slog.Warn("error queueing live runner session event", "runner_id", runner.RunnerID, "session_id", sessionID, "event", event, "err", err) + } +} + +func (runner *liveRunner) closeChannelsLocked() { + for _, channel := range []*liveRunnerTrickleChannel{runner.o2r} { + if channel == nil { + continue + } + if err := channel.close(); err != nil { + slog.Error("error closing live runner bootstrap trickle channel", "channel", channel.channel.ChannelName, "err", err) + } + } +} + +func (runner *liveRunner) liveRunnerSessionLocked(sessionID string) (*liveRunnerSession, error) { + // liveRunnerSessionLocked requires runner.mu. + if runner.removed || !isReadyStatus(runner.Status) { + return nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner not found"} + } + session := runner.sessions[sessionID] + if session == nil { + return nil, &RunnerError{StatusCode: http.StatusNotFound, Message: "runner session not found"} + } + return session, nil +} + +func normalizeTrickleChannelName(sessionID, name string) (string, string, error) { + name = strings.ReplaceAll(strings.TrimSpace(name), "/", "-") + if name == "" { + return "", "", &RunnerError{StatusCode: http.StatusBadRequest, Message: "trickle channel name is required"} + } + if !liveRunnerChannelNamePattern.MatchString(name) { + return "", "", &RunnerError{StatusCode: http.StatusBadRequest, Message: "trickle channel name may only contain alphanumerics, underscores, or hyphens"} + } + return name, sessionID + "-" + name, nil +} + +func (session *liveRunnerSession) closeChannelsLocked() { + for name, channel := range session.channels { + if err := channel.close(); err != nil { + slog.Error("error closing live runner trickle channel", "channel", channel.channel.ChannelName, "err", err) + } + delete(session.channels, name) + } +} + +func writeO2RLoop(channel *liveRunnerTrickleChannel) { + ticker := time.NewTicker(liveRunnerO2RKeepaliveInterval) + defer ticker.Stop() + + for { + select { + case <-channel.done: + return + case msg := <-channel.messages: + if !channel.writeJSON(msg) { + return + } + case <-ticker.C: + if !channel.writeJSON(liveRunnerO2RKeepaliveMessage) { + return + } + } + } +} + +func (channel *liveRunnerTrickleChannel) queueJSON(v any) error { + if channel == nil { + return fmt.Errorf("live runner trickle channel is closed") + } + msg, err := json.Marshal(v) + if err != nil { + return err + } + select { + case <-channel.done: + return fmt.Errorf("live runner trickle channel %q is closed", channel.channel.ChannelName) + case channel.messages <- string(msg): + return nil + default: + return fmt.Errorf("live runner trickle channel %q write queue is full", channel.channel.ChannelName) + } +} + +func (channel *liveRunnerTrickleChannel) writeJSON(msg string) bool { + if err := channel.publisher.Write(strings.NewReader(msg)); err != nil { + if errors.Is(err, trickle.StreamNotFoundErr) { + channel.closeOnce.Do(func() { + close(channel.done) + }) + return false + } + slog.Warn("error writing live runner trickle message", "channel", channel.channel.ChannelName, "err", err) + } + return true +} + +func (channel *liveRunnerTrickleChannel) close() error { + if channel == nil { + return nil + } + channel.closeOnce.Do(func() { + close(channel.done) + }) + return channel.publisher.Close() +} + +func newRunnerSeed() []byte { + return secureRandomBytes(liveRunnerSeedBytes) +} + +func randomStr(n int) string { + return strings.ToLower(liveRunnerEncoding.EncodeToString(secureRandomBytes(n))) +} + +func secureRandomBytes(n int) []byte { + b := make([]byte, n) + if _, err := rand.Read(b); err != nil { + panic(fmt.Sprintf("generate secure random bytes: %v", err)) + } + return b +} + +func (runner *liveRunner) sessionToken(sessionID string) string { + return deriveSecret(runner.seed, "session:"+sessionID) +} + +func (runner *liveRunner) sessionIDsLocked() []string { + // sessionIDsLocked requires runner.mu. + type sessionEntry struct { + id string + createdAt time.Time + } + entries := make([]sessionEntry, 0, len(runner.sessions)) + for id, session := range runner.sessions { + entries = append(entries, sessionEntry{id: id, createdAt: session.createdAt}) + } + sort.Slice(entries, func(i, j int) bool { + if entries[i].createdAt.Equal(entries[j].createdAt) { + return entries[i].id < entries[j].id + } + return entries[i].createdAt.Before(entries[j].createdAt) + }) + sessionIDs := make([]string, 0, len(entries)) + for _, entry := range entries { + sessionIDs = append(sessionIDs, entry.id) + } + return sessionIDs +} + +func deriveSecret(seed []byte, label string) string { + return strings.ToLower(liveRunnerEncoding.EncodeToString(deriveSecretBytes(seed, label))) +} + +func deriveSecretBytes(seed []byte, label string) []byte { + mac := hmac.New(sha256.New, seed) + mac.Write([]byte(label)) + return mac.Sum(nil) +} + +func validSecret(got, want string) bool { + return hmac.Equal([]byte(got), []byte(want)) +} + +func (runner *liveRunner) discoveryRunner() LiveRunnerDiscoveryRunner { + // discoveryRunner requires runner.mu. + priceInfo := runner.PriceInfo + if runner.converter != nil { + converted, err := convertedPriceInfo(runner.converter) + if err != nil { + slog.Error("error reading converted live runner price", "app", runner.App, "endpoint", runner.RunnerURL, "err", err) + } else { + priceInfo = converted + } + } + var discoveryPriceInfo *LiveRunnerPriceInfo + if !runner.offchain && priceInfo != (LiveRunnerPriceInfo{}) { + discoveryPriceInfo = &priceInfo + } + return LiveRunnerDiscoveryRunner{ + GPU: cloneLiveRunnerGPU(runner.GPU), + App: runner.App, + Version: runner.Version, + Mode: runner.Mode, + Capacity: runner.Capacity, + CapacityUsed: len(runner.sessions), + CapacityAvailable: runner.Capacity - len(runner.sessions), + PriceInfo: discoveryPriceInfo, + } +} + +func (runner *liveRunner) updatePriceConverterLocked() { + // updatePriceConverterLocked requires runner.mu. + if runner.offchain { + return + } + // Heartbeats call this each time, but converter rebuild work is only done when + // price_info changes (or when a converter is missing after a prior failure). + if runner.converter != nil && runner.priceSource == runner.PriceInfo { + return + } + runner.stopPriceConverterLocked() + converter, err := newConverterForRunner(runner.PriceInfo) + if err != nil { + slog.Error("error creating live runner price converter", "runner_id", runner.RunnerID, "app", runner.App, "endpoint", runner.RunnerURL, "err", err) + return + } + runner.priceSource = runner.PriceInfo + runner.converter = converter +} + +func (runner *liveRunner) stopPriceConverterLocked() { + // stopPriceConverterLocked requires runner.mu. + if runner.converter == nil { + return + } + runner.converter.Stop() + runner.converter = nil +} + +func newConverterForRunner(priceInfo LiveRunnerPriceInfo) (*core.AutoConvertedPrice, error) { + if strings.ToUpper(priceInfo.Unit) != "USD" { + return nil, nil + } + + usdPerPixel := usdPerPixelFromUSDPerHour(priceInfo) + return core.NewAutoConvertedPrice("USD", usdPerPixel, nil) +} + +func usdPerPixelFromUSDPerHour(priceInfo LiveRunnerPriceInfo) *big.Rat { + pixelsPerHour := 1280 * 720 * 30 * 3600 // 720p @ 30fps + usdPerHour := new(big.Rat).SetFrac64(priceInfo.PricePerUnit, priceInfo.PixelsPerUnit) + return new(big.Rat).Quo(usdPerHour, new(big.Rat).SetInt64(int64(pixelsPerHour))) +} + +func convertedPriceInfo(converter *core.AutoConvertedPrice) (LiveRunnerPriceInfo, error) { + if converter == nil { + return LiveRunnerPriceInfo{}, nil + } + price, err := common.PriceToInt64(converter.Value()) + if err != nil { + return LiveRunnerPriceInfo{}, err + } + return LiveRunnerPriceInfo{ + PricePerUnit: price.Num().Int64(), + PixelsPerUnit: price.Denom().Int64(), + Unit: "WEI", + }, nil +} + +func isReadyStatus(status string) bool { + return status == "" || status == "ready" +} + +func cloneLiveRunnerGPU(gpu *LiveRunnerGPU) *LiveRunnerGPU { + if gpu == nil { + return nil + } + copy := *gpu + return © +} diff --git a/ai/runner/live_runner_test.go b/ai/runner/live_runner_test.go new file mode 100644 index 0000000000..a34f2c4064 --- /dev/null +++ b/ai/runner/live_runner_test.go @@ -0,0 +1,1744 @@ +package runner + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "math/big" + "net/http" + "net/http/httptest" + "net/url" + "regexp" + "slices" + "strings" + "sync" + "testing" + "time" + + "github.com/livepeer/go-livepeer/common" + "github.com/livepeer/go-livepeer/core" + "github.com/livepeer/go-livepeer/eth" + "github.com/livepeer/go-livepeer/trickle" +) + +type stubPriceFeedWatcher struct { + price eth.PriceData +} + +func (s stubPriceFeedWatcher) Currencies() (string, string, error) { + return "ETH", "USD", nil +} + +func (s stubPriceFeedWatcher) Current() (eth.PriceData, error) { + return s.price, nil +} + +func (s stubPriceFeedWatcher) Subscribe(context.Context, chan<- eth.PriceData) {} + +func liveRunnerTestHeartbeat(runnerID string) LiveRunnerHeartbeatRequest { + return LiveRunnerHeartbeatRequest{ + RunnerID: runnerID, + Label: "test-runner", + RunnerURL: "https://runner.example.com", + Version: "1.2.3", + Status: "ready", + GPU: &LiveRunnerGPU{ + ID: "gpu-0", + Name: "NVIDIA L40S", + VRAMMB: 46068, + }, + App: "new-ai-pipeline/model-a", + Capacity: 1, + } +} + +const liveRunnerTestBootstrapSecret = "bootstrap-secret" + +var ( + liveRunnerTestGeneratedRunnerIDPattern = regexp.MustCompile(`^runner_[a-z2-7]{8}$`) + liveRunnerTestGeneratedSessionIDPattern = regexp.MustCompile(`^session_[a-z2-7]{8}$`) + liveRunnerTestSecretPattern = regexp.MustCompile(`^[a-z2-7]+$`) +) + +type liveRunnerTestHost struct{} + +func (liveRunnerTestHost) ServiceURI() *url.URL { + u, _ := url.Parse("http://localhost:1234") + return u +} + +func (h liveRunnerTestHost) LiveRunnerURI() *url.URL { + return h.ServiceURI() +} + +func (liveRunnerTestHost) RegistrationSecret() string { + return liveRunnerTestBootstrapSecret +} + +type liveRunnerSplitURIHost struct{} + +func (liveRunnerSplitURIHost) ServiceURI() *url.URL { + u, _ := url.Parse("https://public.example.com") + return u +} + +func (liveRunnerSplitURIHost) LiveRunnerURI() *url.URL { + u, _ := url.Parse("http://go-livepeer:8935") + return u +} + +func (liveRunnerSplitURIHost) RegistrationSecret() string { + return liveRunnerTestBootstrapSecret +} + +type liveRunnerTestHostWithoutSecret struct{} + +func (liveRunnerTestHostWithoutSecret) ServiceURI() *url.URL { + u, _ := url.Parse("http://localhost:1234") + return u +} + +func (h liveRunnerTestHostWithoutSecret) LiveRunnerURI() *url.URL { + return h.ServiceURI() +} + +func (liveRunnerTestHostWithoutSecret) RegistrationSecret() string { + return "" +} + +func newLiveRunnerTestRegistry() *LiveRunnerRegistry { + return NewLiveRunnerRegistry(LiveRunnerRegistryConfig{Host: liveRunnerTestHost{}}) +} + +func newOnchainLiveRunnerTestRegistry() *LiveRunnerRegistry { + return NewLiveRunnerRegistry(LiveRunnerRegistryConfig{Host: liveRunnerTestHost{}, Onchain: true}) +} + +func liveRunnerTestRegister(t *testing.T, registry *LiveRunnerRegistry, req LiveRunnerHeartbeatRequest) *LiveRunnerHeartbeatResponse { + t.Helper() + ensureLiveRunnerTestTrickleServer(t, registry) + resp, err := registry.Heartbeat(req, liveRunnerTestBootstrapSecret) + if err != nil { + t.Fatal(err) + } + if resp.HeartbeatSecret == "" { + t.Fatal("expected heartbeat secret") + } + if !liveRunnerTestSecretPattern.MatchString(resp.HeartbeatSecret) { + t.Fatalf("expected base32 heartbeat secret, got %q", resp.HeartbeatSecret) + } + if resp.Orchestrator != "http://localhost:1234" { + t.Fatalf("unexpected orchestrator URL: %s", resp.Orchestrator) + } + return resp +} + +func ensureLiveRunnerTestTrickleServer(t *testing.T, registry *LiveRunnerRegistry) (string, func(string) int) { + t.Helper() + registry.mu.Lock() + hasTrickleServer := registry.trickleSrv != nil + channelBaseURL := registry.internalTrickleBaseURL + registry.mu.Unlock() + if hasTrickleServer { + return channelBaseURL, nil + } + trickleSrv, channelBaseURL, channelStatus := newLiveRunnerTestTrickleServer(t) + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + return channelBaseURL, channelStatus +} + +func liveRunnerTestBootstrapChannels(resp *LiveRunnerHeartbeatResponse) []LiveRunnerTrickleChannel { + channels := []LiveRunnerTrickleChannel{} + if resp.O2R != nil { + channels = append(channels, *resp.O2R) + } + return channels +} + +func TestLiveRunnerRegistry_HeartbeatUpsertCapacity(t *testing.T) { + registry := newLiveRunnerTestRegistry() + + resp := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("")) + if resp.RunnerID == "" { + t.Fatal("expected runner id") + } + if !liveRunnerTestGeneratedRunnerIDPattern.MatchString(resp.RunnerID) { + t.Fatalf("expected generated base32 runner id, got %q", resp.RunnerID) + } + runner := registry.runners[resp.RunnerID] + if runner == nil || runner.Capacity != 1 { + t.Fatalf("unexpected initial runner state: %+v", runner) + } + + _, err := registry.Heartbeat(LiveRunnerHeartbeatRequest{ + RunnerID: resp.RunnerID, + Label: "test-runner", + RunnerURL: "https://runner.example.com", + Status: "ready", + App: "new-ai-pipeline/model-a", + Capacity: 2, + }, resp.HeartbeatSecret) + if err != nil { + t.Fatal(err) + } + runner = registry.runners[resp.RunnerID] + if runner == nil || runner.Capacity != 2 { + t.Fatalf("unexpected heartbeat runner state: %+v", runner) + } +} + +func TestLiveRunnerRegistry_SplitsPublicAndRunnerURIs(t *testing.T) { + registry := NewLiveRunnerRegistry(LiveRunnerRegistryConfig{Host: liveRunnerSplitURIHost{}}) + t.Cleanup(registry.Stop) + trickleSrv, _, _ := newLiveRunnerTestTrickleServer(t) + registry.SetTrickleServer(trickleSrv, "https://public.example.com/ai/trickle/", "http://go-livepeer:8935/ai/trickle/") + + resp, err := registry.Heartbeat(liveRunnerTestHeartbeat("runner-split-uri"), liveRunnerTestBootstrapSecret) + if err != nil { + t.Fatal(err) + } + if resp.Orchestrator != "http://go-livepeer:8935" { + t.Fatalf("expected heartbeat orchestrator to use runner URI, got %q", resp.Orchestrator) + } + if resp.O2R == nil || !strings.HasPrefix(resp.O2R.URL, "http://go-livepeer:8935/ai/trickle/") { + t.Fatalf("expected O2R channel to use runner trickle URL, got %+v", resp.O2R) + } + if resp.O2R.InternalURL != "" { + t.Fatalf("expected O2R channel to omit internal URL, got %+v", resp.O2R) + } + + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one discovery runner, got %d", len(runners)) + } + if runners[0].URL != "https://public.example.com/apps/runner-split-uri/session" { + t.Fatalf("expected discovery URL to use public service URI, got %q", runners[0].URL) + } + + sessionID, _, err := registry.ReserveSession("runner-split-uri") + if err != nil { + t.Fatal(err) + } + channel, err := registry.CreateTrickleChannel("runner-split-uri", sessionID, "events", "application/json") + if err != nil { + t.Fatal(err) + } + if !strings.HasPrefix(channel.URL, "https://public.example.com/ai/trickle/") { + t.Fatalf("expected runner-created trickle channel to use public trickle URL, got %q", channel.URL) + } + if !strings.HasPrefix(channel.InternalURL, "http://go-livepeer:8935/ai/trickle/") { + t.Fatalf("expected runner-created trickle channel to include internal trickle URL, got %q", channel.InternalURL) + } + token, err := registry.SessionTokenForSession("runner-split-uri", sessionID) + if err != nil { + t.Fatal(err) + } + if err := registry.ValidSessionToken("runner-split-uri", sessionID, token); err != nil { + t.Fatal(err) + } +} + +func TestLiveRunnerRegistry_HeartbeatUnknownIDCreatesRunner(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner_custom_1") + + resp := liveRunnerTestRegister(t, registry, req) + if resp.RunnerID != "runner_custom_1" { + t.Fatalf("expected supplied runner id to be used, got %s", resp.RunnerID) + } + runner := registry.runners[resp.RunnerID] + if runner == nil || runner.Capacity != 1 { + t.Fatalf("unexpected state after create-by-unknown-id heartbeat: %+v", runner) + } +} + +func TestLiveRunnerRegistry_HeartbeatCreatesBootstrapTrickleChannels(t *testing.T) { + trickleSrv, channelBaseURL, channelStatus := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + req := liveRunnerTestHeartbeat("runner-bootstrap") + + resp := liveRunnerTestRegister(t, registry, req) + channels := liveRunnerTestBootstrapChannels(resp) + if len(channels) != 1 { + t.Fatalf("expected bootstrap trickle channels, got %+v", resp) + } + expectedNames := []string{ + LiveRunnerTrickleOrchestratorToRunner, + } + for i, expectedName := range expectedNames { + channel := channels[i] + if channel.Name != expectedName { + t.Fatalf("unexpected bootstrap channel name at %d: %+v", i, channel) + } + if !strings.HasPrefix(channel.ChannelName, resp.RunnerID+"-") || !strings.HasSuffix(channel.ChannelName, "-"+expectedName) { + t.Fatalf("expected randomized backing channel for %q, got %q", expectedName, channel.ChannelName) + } + if channel.ChannelName == resp.RunnerID+"-"+expectedName { + t.Fatalf("expected non-guessable backing channel name, got %q", channel.ChannelName) + } + if channel.URL != channelBaseURL+channel.ChannelName { + t.Fatalf("unexpected bootstrap channel URL: %s", channel.URL) + } + if channel.MimeType != "application/json" { + t.Fatalf("unexpected bootstrap channel MIME type: %s", channel.MimeType) + } + if status := channelStatus(channel.ChannelName); status != http.StatusOK { + t.Fatalf("expected bootstrap channel to exist, got status=%d", status) + } + } +} + +func TestLiveRunnerRegistry_O2RKeepaliveWriteLoop(t *testing.T) { + restore := liveRunnerTestSetO2RKeepaliveInterval(t, 10*time.Millisecond) + defer restore() + + trickleSrv, channelBaseURL, _ := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + t.Cleanup(registry.Stop) + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + resp := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-o2r-keepalive")) + + msg := liveRunnerTestReadTrickleMessage(t, trickleSrv, resp.O2R.ChannelName, func(msg string) bool { + return msg == liveRunnerO2RKeepaliveMessage + }) + if msg != liveRunnerO2RKeepaliveMessage { + t.Fatalf("unexpected keepalive message: %q", msg) + } +} + +func TestLiveRunnerRegistry_O2RSessionLifecycleMessages(t *testing.T) { + restore := liveRunnerTestSetO2RKeepaliveInterval(t, time.Hour) + defer restore() + + trickleSrv, channelBaseURL, _ := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + t.Cleanup(registry.Stop) + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + req := liveRunnerTestHeartbeat("runner-o2r-session") + req.Capacity = 2 + resp := liveRunnerTestRegister(t, registry, req) + + sessionID, _, err := registry.ReserveSession("runner-o2r-session", "manifest-id") + if err != nil { + t.Fatal(err) + } + liveRunnerTestReadSessionEvent(t, trickleSrv, resp.O2R.ChannelName, "reserved", sessionID) + + if err := registry.ReleaseSession("runner-o2r-session", sessionID); err != nil { + t.Fatal(err) + } + liveRunnerTestReadSessionEvent(t, trickleSrv, resp.O2R.ChannelName, "released", sessionID) +} + +func TestLiveRunnerRegistry_O2RWriteLoopClosesOnUnregister(t *testing.T) { + restore := liveRunnerTestSetO2RKeepaliveInterval(t, time.Hour) + defer restore() + + registry := newLiveRunnerTestRegistry() + t.Cleanup(registry.Stop) + resp := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-o2r-close")) + + registry.mu.Lock() + runner := registry.runners["runner-o2r-close"] + registry.mu.Unlock() + if runner == nil || runner.o2r == nil { + t.Fatal("expected registered runner with o2r channel") + } + o2r := runner.o2r + + if err := registry.Unregister("runner-o2r-close", resp.HeartbeatSecret); err != nil { + t.Fatal(err) + } + select { + case <-o2r.done: + case <-time.After(time.Second): + t.Fatal("expected o2r write loop to close on unregister") + } +} + +func TestLiveRunnerRegistry_HeartbeatDoesNotReturnBootstrapChannelsAgain(t *testing.T) { + registry := newLiveRunnerTestRegistry() + resp := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-bootstrap-once")) + + nextResp, err := registry.Heartbeat(liveRunnerTestHeartbeat(resp.RunnerID), resp.HeartbeatSecret) + if err != nil { + t.Fatal(err) + } + if nextResp.O2R != nil { + t.Fatalf("expected no bootstrap channels on follow-up heartbeat, got %+v", nextResp) + } + registry.mu.Lock() + runner := registry.runners[resp.RunnerID] + registry.mu.Unlock() + runner.mu.Lock() + defer runner.mu.Unlock() + if runner.o2r == nil { + t.Fatalf("expected existing bootstrap channel to be preserved, got o2r=%+v", runner.o2r) + } +} + +func TestLiveRunnerRegistry_HeartbeatSessionIDs(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-session-ids") + req.Capacity = 3 + req.SessionIDs = []string{"runner-reported-session"} + resp := liveRunnerTestRegister(t, registry, req) + if len(resp.SessionIDs) != 0 { + t.Fatalf("expected no active session ids on initial heartbeat, got %+v", resp.SessionIDs) + } + + registry.mu.Lock() + runner := registry.runners[resp.RunnerID] + registry.mu.Unlock() + runner.mu.Lock() + if !slices.Equal(runner.LiveRunnerHeartbeatRequest.SessionIDs, []string{"runner-reported-session"}) { + t.Fatalf("expected request session_ids to be retained, got %+v", runner.LiveRunnerHeartbeatRequest.SessionIDs) + } + runner.mu.Unlock() + + if _, _, err := registry.ReserveSession(resp.RunnerID, "session-z-oldest"); err != nil { + t.Fatal(err) + } + time.Sleep(time.Millisecond) + if _, _, err := registry.ReserveSession(resp.RunnerID, "session-a-newest"); err != nil { + t.Fatal(err) + } + + followupReq := liveRunnerTestHeartbeat(resp.RunnerID) + followupReq.Capacity = 3 + followupReq.SessionIDs = []string{"runner-reported-session"} + nextResp, err := registry.Heartbeat(followupReq, resp.HeartbeatSecret) + if err != nil { + t.Fatal(err) + } + want := []string{"session-z-oldest", "session-a-newest"} + if !slices.Equal(nextResp.SessionIDs, want) { + t.Fatalf("unexpected session ids: got %+v want %+v", nextResp.SessionIDs, want) + } + + registry.mu.Lock() + runner = registry.runners[resp.RunnerID] + registry.mu.Unlock() + runner.mu.Lock() + defer runner.mu.Unlock() + if !slices.Equal(runner.LiveRunnerHeartbeatRequest.SessionIDs, []string{"runner-reported-session"}) { + t.Fatalf("expected follow-up request session_ids to be retained, got %+v", runner.LiveRunnerHeartbeatRequest.SessionIDs) + } +} + +func TestLiveRunnerRegistry_InvalidHeartbeatAuthDoesNotLeakRegistryLock(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-lock-leak") + resp := liveRunnerTestRegister(t, registry, req) + + const attempts = 200 + done := make(chan error, attempts+1) + start := make(chan struct{}) + + for i := 0; i < attempts; i++ { + go func(i int) { + <-start + badReq := req + _, err := registry.Heartbeat(badReq, fmt.Sprintf("wrong-secret-%d", i)) + if !isRunnerErrorStatus(err, http.StatusUnauthorized) { + done <- fmt.Errorf("wrong-secret heartbeat %d: expected unauthorized, got %v", i, err) + return + } + done <- nil + }(i) + } + + go func() { + <-start + for i := 0; i < attempts; i++ { + if runners := registry.Runners(); len(runners) != 1 { + done <- fmt.Errorf("Runners during bad-auth storm: expected one runner, got %d", len(runners)) + return + } + goodReq := req + if _, err := registry.Heartbeat(goodReq, resp.HeartbeatSecret); err != nil { + done <- fmt.Errorf("valid heartbeat during bad-auth storm: %w", err) + return + } + } + done <- nil + }() + + close(start) + timeout := time.After(5 * time.Second) + for i := 0; i < attempts+1; i++ { + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-timeout: + t.Fatal("timed out waiting for concurrent heartbeat auth checks; registry lock may be leaked") + } + } +} + +func TestLiveRunnerRegistry_HeartbeatDisabledWithoutRegistrationSecret(t *testing.T) { + registry := NewLiveRunnerRegistry(LiveRunnerRegistryConfig{Host: liveRunnerTestHostWithoutSecret{}}) + _, err := registry.Heartbeat(liveRunnerTestHeartbeat("runner-1"), "") + if !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected dynamic registration disabled error, got %v", err) + } +} + +func TestLiveRunnerRegistry_DefaultCapacityWhenUnset(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner_default_capacity") + req.Capacity = 0 + + resp := liveRunnerTestRegister(t, registry, req) + registry.mu.Lock() + runner := registry.runners[resp.RunnerID] + registry.mu.Unlock() + if runner == nil || runner.Capacity != 1 { + t.Fatalf("expected default capacity=1, got %+v", runner) + } +} + +func TestLiveRunnerRegistry_DefaultModeWhenUnset(t *testing.T) { + registry := newLiveRunnerTestRegistry() + resp := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner_default_mode")) + + mode, err := registry.RunnerMode(resp.RunnerID) + if err != nil { + t.Fatal(err) + } + if mode != LiveRunnerModePersistent { + t.Fatalf("expected default mode %q, got %q", LiveRunnerModePersistent, mode) + } +} + +func TestLiveRunnerRegistry_HeartbeatSingleShotMode(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-single-shot") + req.Mode = LiveRunnerModeSingleShot + resp := liveRunnerTestRegister(t, registry, req) + + mode, err := registry.RunnerMode(resp.RunnerID) + if err != nil { + t.Fatal(err) + } + if mode != LiveRunnerModeSingleShot { + t.Fatalf("expected mode %q, got %q", LiveRunnerModeSingleShot, mode) + } +} + +func TestLiveRunnerRegistry_HeartbeatSingleShotModeAlias(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-single-shot-alias") + req.Mode = "single_shot" + resp := liveRunnerTestRegister(t, registry, req) + + mode, err := registry.RunnerMode(resp.RunnerID) + if err != nil { + t.Fatal(err) + } + if mode != LiveRunnerModeSingleShot { + t.Fatalf("expected mode %q, got %q", LiveRunnerModeSingleShot, mode) + } +} + +func TestLiveRunnerRegistry_InvalidMode(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-invalid-mode") + req.Mode = "invalid" + if _, err := registry.Heartbeat(req, liveRunnerTestBootstrapSecret); !isRunnerErrorStatus(err, http.StatusBadRequest) || !strings.Contains(err.Error(), "mode") { + t.Fatalf("expected invalid mode bad request, got %v", err) + } +} + +func TestLiveRunnerRegistry_OnchainRequiresHeartbeatPriceInfo(t *testing.T) { + registry := newOnchainLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-missing-price") + + _, err := registry.Heartbeat(req, liveRunnerTestBootstrapSecret) + if !isRunnerErrorStatus(err, http.StatusBadRequest) || !strings.Contains(err.Error(), "price_info") { + t.Fatalf("expected missing price_info bad request, got %v", err) + } +} + +func TestLiveRunnerRegistry_OnchainRequiresStaticRunnerPriceInfo(t *testing.T) { + registry := newOnchainLiveRunnerTestRegistry() + + _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "static-missing-price", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + HealthURL: "https://runner.example.com/health", + }}}) + if !isRunnerErrorStatus(err, http.StatusBadRequest) || !strings.Contains(err.Error(), "price_info") { + t.Fatalf("expected missing price_info bad request, got %v", err) + } +} + +func TestLiveRunnerRegistry_OnchainAcceptsWEIPrice(t *testing.T) { + registry := newOnchainLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-wei-price") + req.PriceInfo = LiveRunnerPriceInfo{PricePerUnit: 10, PixelsPerUnit: 2, Unit: "wei"} + + liveRunnerTestRegister(t, registry, req) + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one runner, got %d", len(runners)) + } + if runners[0].PriceInfo == nil || *runners[0].PriceInfo != (LiveRunnerPriceInfo{PricePerUnit: 10, PixelsPerUnit: 2, Unit: "WEI"}) { + t.Fatalf("unexpected runner price info: %+v", runners[0].PriceInfo) + } +} + +func TestLiveRunnerRegistry_OffchainIgnoresUSDPriceWatcher(t *testing.T) { + prevWatcher := core.PriceFeedWatcher + core.PriceFeedWatcher = nil + defer func() { core.PriceFeedWatcher = prevWatcher }() + + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-offchain-usd") + req.PriceInfo = LiveRunnerPriceInfo{PricePerUnit: 10, PixelsPerUnit: 1, Unit: "USD"} + + liveRunnerTestRegister(t, registry, req) + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one runner, got %d", len(runners)) + } + if runners[0].PriceInfo != nil { + t.Fatalf("expected offchain discovery price to be suppressed, got %+v", runners[0].PriceInfo) + } +} + +func TestLiveRunnerRegistry_ClonesRunnerGPUMetadata(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-cloned-gpu") + req.GPU.Name = "original-dynamic" + liveRunnerTestRegister(t, registry, req) + req.GPU.Name = "mutated-dynamic" + + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one dynamic runner, got %d", len(runners)) + } + if runners[0].GPU == nil || runners[0].GPU.Name != "original-dynamic" { + t.Fatalf("expected dynamic runner GPU to be cloned, got %+v", runners[0].GPU) + } + + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + staticGPU := &LiveRunnerGPU{Name: "original-static"} + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "static-cloned-gpu", + RunnerURL: "https://runner.example.com", + GPU: staticGPU, + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + }}}); err != nil { + t.Fatal(err) + } + staticGPU.Name = "mutated-static" + + runners = registry.Runners() + var foundStatic bool + for _, runner := range runners { + if runner.App != "live-video-to-video/scope" { + continue + } + foundStatic = true + if runner.GPU == nil || runner.GPU.Name != "original-static" { + t.Fatalf("expected static runner GPU to be cloned, got %+v", runner.GPU) + } + } + if !foundStatic { + t.Fatal("expected static runner in discovery") + } +} + +func TestParseStaticLiveRunnerConfigDefaultsHealthStatusAndNumericGPU(t *testing.T) { + cfg, err := ParseStaticLiveRunnerConfig([]byte(`{"runners":[{"label":"app","runner_url":"https://runner.example.com","app":"live-video-to-video/scope","capacity":1,"health_url":"https://runner.example.com/health","gpu":-1}]}`)) + if err != nil { + t.Fatal(err) + } + if len(cfg.Runners) != 1 { + t.Fatalf("expected one runner, got %d", len(cfg.Runners)) + } + if cfg.Runners[0].Route != LiveRunnerRoutingRunnerID { + t.Fatalf("expected default routing %q, got %q", LiveRunnerRoutingRunnerID, cfg.Runners[0].Route) + } + if cfg.Runners[0].HealthyStatusCode != http.StatusOK { + t.Fatalf("expected default health status 200, got %d", cfg.Runners[0].HealthyStatusCode) + } + if cfg.Runners[0].GPU == nil || cfg.Runners[0].GPU.ID != "-1" { + t.Fatalf("expected numeric gpu fallback, got %+v", cfg.Runners[0].GPU) + } + if cfg.Runners[0].Mode != "" { + t.Fatalf("expected raw static config mode to remain empty before registration, got %q", cfg.Runners[0].Mode) + } +} + +func TestParseStaticLiveRunnerConfigRoute(t *testing.T) { + cfg, err := ParseStaticLiveRunnerConfig([]byte(`{"runners":[{"label":"app","routing":"label","runner_url":"https://runner.example.com","app":"live-video-to-video/scope","health_url":"https://runner.example.com/health"}]}`)) + if err != nil { + t.Fatal(err) + } + if cfg.Runners[0].Route != LiveRunnerRoutingLabel { + t.Fatalf("expected routing %q, got %q", LiveRunnerRoutingLabel, cfg.Runners[0].Route) + } + if _, err := ParseStaticLiveRunnerConfig([]byte(`{"runners":[{"label":"app","routing":"bad","runner_url":"https://runner.example.com","app":"live-video-to-video/scope","health_url":"https://runner.example.com/health"}]}`)); err == nil || !strings.Contains(err.Error(), "routing") { + t.Fatalf("expected invalid routing error, got %v", err) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersResolvesHealthPath(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/health" { + w.WriteHeader(http.StatusNotFound) + return + } + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + resp, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "path-health", + RunnerURL: healthSrv.URL, + App: "live-video-to-video/scope", + Capacity: 1, + HealthURL: "/health", + }}}) + if err != nil { + t.Fatal(err) + } + if len(resp.Runners) != 1 { + t.Fatalf("expected one runner, got %d", len(resp.Runners)) + } + if want := healthSrv.URL + "/health"; resp.Runners[0].HealthURL != want { + t.Fatalf("expected health URL %q, got %q", want, resp.Runners[0].HealthURL) + } + if !resp.Runners[0].Healthy { + t.Fatal("expected health path to be checked against runner URL") + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersSingleShotMode(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + resp, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "static-single-shot", + Route: LiveRunnerRoutingLabel, + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Mode: LiveRunnerModeSingleShot, + Capacity: 1, + HealthURL: healthSrv.URL, + }}}) + if err != nil { + t.Fatal(err) + } + if resp.Runners[0].Mode != LiveRunnerModeSingleShot { + t.Fatalf("expected registration mode %q, got %q", LiveRunnerModeSingleShot, resp.Runners[0].Mode) + } + mode, err := registry.RunnerMode("static-single-shot") + if err != nil { + t.Fatal(err) + } + if mode != LiveRunnerModeSingleShot { + t.Fatalf("expected runner mode %q, got %q", LiveRunnerModeSingleShot, mode) + } + runners := registry.Runners() + if len(runners) != 1 || runners[0].URL != "http://localhost:1234/apps/static-single-shot/app" { + t.Fatalf("unexpected single-shot route discovery: %+v", runners) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersHealthAndNoHeartbeatExpiry(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + resp, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "static-app", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 1, + HealthURL: healthSrv.URL, + }}}) + if err != nil { + t.Fatal(err) + } + if len(resp.Runners) != 1 || resp.Runners[0].RunnerID == "" || !resp.Runners[0].Healthy { + t.Fatalf("unexpected static registration response: %+v", resp) + } + runnerID := resp.Runners[0].RunnerID + if _, _, err := registry.ReserveSession(runnerID); err != nil { + t.Fatalf("expected healthy static runner to reserve session: %v", err) + } + + registry.heartbeatTTL = time.Nanosecond + time.Sleep(time.Millisecond) + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected static runner to skip heartbeat expiry, got %d", len(runners)) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersAtomicValidation(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + valid := StaticLiveRunnerConfigEntry{ + Label: "runner-a", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + } + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{valid}}); err != nil { + t.Fatal(err) + } + if len(registry.Runners()) != 1 { + t.Fatal("expected initial static runner") + } + + invalid := valid + invalid.Label = "runner-b" + invalid.RunnerURL = "bad url" + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{valid, invalid}}); !isRunnerErrorStatus(err, http.StatusBadRequest) { + t.Fatalf("expected bad request for invalid batch, got %v", err) + } + if len(registry.Runners()) != 1 { + t.Fatalf("expected invalid batch to leave existing runners unchanged, got %d", len(registry.Runners())) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersRequiresUniqueLabels(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + entry := StaticLiveRunnerConfigEntry{ + Label: "runner-a", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + } + duplicate := entry + duplicate.RunnerURL = "https://runner-two.example.com" + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{entry, duplicate}}); !isRunnerErrorStatus(err, http.StatusBadRequest) { + t.Fatalf("expected duplicate label error, got %v", err) + } + + missingLabel := entry + missingLabel.Label = "" + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{missingLabel}}); !isRunnerErrorStatus(err, http.StatusBadRequest) || !strings.Contains(err.Error(), "label") { + t.Fatalf("expected missing label error, got %v", err) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersUpsertsWithoutDroppingSession(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + entry := StaticLiveRunnerConfigEntry{ + Label: "first", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 2, + HealthURL: healthSrv.URL, + } + resp1, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{entry}}) + if err != nil { + t.Fatal(err) + } + runnerID := resp1.Runners[0].RunnerID + sessionID, _, err := registry.ReserveSession(runnerID) + if err != nil { + t.Fatal(err) + } + + entry.RunnerURL = "https://runner-updated.example.com" + resp2, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{entry}}) + if err != nil { + t.Fatal(err) + } + if resp2.Runners[0].RunnerID != runnerID { + t.Fatalf("expected static runner ID to be preserved, got %s want %s", resp2.Runners[0].RunnerID, runnerID) + } + if _, err := registry.RunnerEndpointForSession(runnerID, sessionID); err != nil { + t.Fatalf("expected active session to survive healthy static upsert: %v", err) + } +} + +func TestLiveRunnerRegistry_RegisterStaticRunnersConcurrentUpserts(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + defer registry.Stop() + + entryA := StaticLiveRunnerConfigEntry{ + Label: "runner-a", + RunnerURL: "https://runner-a.example.com", + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + } + entryB := StaticLiveRunnerConfigEntry{ + Label: "runner-b", + RunnerURL: "https://runner-b.example.com", + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + } + + const registrationWorkers = 100 + var wg sync.WaitGroup + errCh := make(chan error, registrationWorkers) + for i := 0; i < registrationWorkers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + entries := []StaticLiveRunnerConfigEntry{entryA, entryB} + if i%2 == 1 { + entries = []StaticLiveRunnerConfigEntry{entryB, entryA} + } + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: entries}); err != nil { + errCh <- err + } + }(i) + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("concurrent static registrations timed out") + } + close(errCh) + for err := range errCh { + t.Fatal(err) + } + if runners := registry.Runners(); len(runners) != 2 { + t.Fatalf("expected concurrent static upserts to settle on two runners, got %d", len(runners)) + } +} + +func TestLiveRunnerRegistry_StaticRunnerRouteLabel(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{ + Runners: []StaticLiveRunnerConfigEntry{{ + Label: "first", + Route: LiveRunnerRoutingLabel, + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 2, + HealthURL: healthSrv.URL, + }}, + }) + if err != nil { + t.Fatal(err) + } + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one runner, got %d", len(runners)) + } + if runners[0].URL != "http://localhost:1234/apps/first/session" { + t.Fatalf("unexpected label route url: %s", runners[0].URL) + } + sessionID, _, err := registry.ReserveSession("first") + if err != nil { + t.Fatalf("expected label route to reserve: %v", err) + } + if _, err := registry.RunnerEndpointForSession("first", sessionID); err != nil { + t.Fatalf("expected label route to resolve: %v", err) + } +} + +func TestLiveRunnerRegistry_StaticRunnerRouteChangesOnUpsert(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + entry := StaticLiveRunnerConfigEntry{ + Label: "stable-route", + Route: LiveRunnerRoutingRunnerID, + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + HealthURL: healthSrv.URL, + } + resp, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{entry}}) + if err != nil { + t.Fatal(err) + } + runnerID := resp.Runners[0].RunnerID + + entry.Route = LiveRunnerRoutingLabel + if _, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{entry}}); err != nil { + t.Fatal(err) + } + if _, _, err := registry.ReserveSession("stable-route"); err != nil { + t.Fatalf("expected upsert to move static runner to label route: %v", err) + } + if _, _, err := registry.ReserveSession(runnerID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected original runner-id route to be removed, got %v", err) + } +} + +func TestLiveRunnerRegistry_StaticRunnerCustomHealthStatusAndUnhealthyRelease(t *testing.T) { + statusCode := http.StatusCreated + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(statusCode) + })) + defer healthSrv.Close() + + registry := newLiveRunnerTestRegistry() + resp, err := registry.RegisterStaticRunners(StaticLiveRunnerConfig{Runners: []StaticLiveRunnerConfigEntry{{ + Label: "custom-health", + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 1, + HealthURL: healthSrv.URL, + HealthyStatusCode: http.StatusCreated, + }}}) + if err != nil { + t.Fatal(err) + } + runnerID := resp.Runners[0].RunnerID + sessionID, _, err := registry.ReserveSession(runnerID) + if err != nil { + t.Fatal(err) + } + statusCode = http.StatusOK + registry.checkStaticRunnerHealth() + if _, err := registry.RunnerEndpointForSession(runnerID, sessionID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected unhealthy static runner to release sessions, got %v", err) + } + if runners := registry.Runners(); len(runners) != 0 { + data, _ := json.Marshal(runners) + t.Fatalf("expected unhealthy static runner to be hidden, got %s", data) + } +} + +func TestLiveRunnerRegistry_RunnersDiscoveryShape(t *testing.T) { + registry := newLiveRunnerTestRegistry() + + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner_discovery_1")) + + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one runner, got %d", len(runners)) + } + runner := runners[0] + if runner.URL != "http://localhost:1234/apps/runner_discovery_1/session" { + t.Fatalf("unexpected url: %s", runner.URL) + } + if runner.GPU == nil || runner.GPU.Name != "NVIDIA L40S" { + t.Fatalf("unexpected gpu: %+v", runner.GPU) + } + if runner.App != "new-ai-pipeline/model-a" { + t.Fatalf("unexpected app: %s", runner.App) + } + if runner.PriceInfo != nil { + t.Fatalf("unexpected runner price info: %+v", runner.PriceInfo) + } + if runner.Version != "1.2.3" { + t.Fatalf("unexpected runner version: %s", runner.Version) + } +} + +func TestLiveRunnerRegistry_ConvertsUSDToWEI(t *testing.T) { + prevWatcher := core.PriceFeedWatcher + core.PriceFeedWatcher = stubPriceFeedWatcher{price: eth.PriceData{Price: big.NewRat(2000, 1)}} + defer func() { core.PriceFeedWatcher = prevWatcher }() + + registry := newOnchainLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-1") + req.PriceInfo = LiveRunnerPriceInfo{PricePerUnit: 10, PixelsPerUnit: 1, Unit: "USD"} + liveRunnerTestRegister(t, registry, req) + + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one runner, got %d", len(runners)) + } + if runners[0].PriceInfo == nil { + t.Fatal("expected runner price info") + } + got := big.NewRat(runners[0].PriceInfo.PricePerUnit, runners[0].PriceInfo.PixelsPerUnit) + expected, err := common.PriceToInt64(big.NewRat(5_000_000_000_000_000, 1280*720*30*3600)) + if err != nil { + t.Fatal(err) + } + if got.Cmp(expected) != 0 { + t.Fatalf("unexpected converted price: got=%s want=%s", got, expected) + } + if runners[0].PriceInfo.Unit != "WEI" { + t.Fatalf("unexpected converted unit: %s", runners[0].PriceInfo.Unit) + } +} + +func TestLiveRunnerRegistry_SharedURLAppKeepsPerRunnerPrices(t *testing.T) { + registry := newOnchainLiveRunnerTestRegistry() + + req1 := liveRunnerTestHeartbeat("runner-1") + req1.PriceInfo = LiveRunnerPriceInfo{PricePerUnit: 10, PixelsPerUnit: 1, Unit: "WEI"} + req2 := liveRunnerTestHeartbeat("runner-2") + req2.PriceInfo = LiveRunnerPriceInfo{PricePerUnit: 20, PixelsPerUnit: 1, Unit: "WEI"} + resp1 := liveRunnerTestRegister(t, registry, req1) + liveRunnerTestRegister(t, registry, req2) + + if err := registry.Unregister("runner-1", resp1.HeartbeatSecret); err != nil { + t.Fatal(err) + } + runners := registry.Runners() + if len(runners) != 1 { + t.Fatalf("expected one remaining runner, got %d", len(runners)) + } + if runners[0].PriceInfo == nil { + t.Fatal("expected runner price info") + } + if got := runners[0].PriceInfo.PricePerUnit; got != 20 { + t.Fatalf("expected remaining runner price to stay isolated, got %d", got) + } +} + +func TestLiveRunnerRegistry_ReserveSessionCapacity(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-1") + req.Capacity = 2 + liveRunnerTestRegister(t, registry, req) + + sessionID1, endpoint, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + if endpoint != req.RunnerURL { + t.Fatalf("unexpected endpoint: %s", endpoint) + } + if !liveRunnerTestGeneratedSessionIDPattern.MatchString(sessionID1) { + t.Fatalf("expected generated base32 session id, got %q", sessionID1) + } + sessionID2, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + if sessionID1 == sessionID2 { + t.Fatalf("expected unique session IDs, got %s", sessionID1) + } + + registry.mu.Lock() + active := len(registry.runners["runner-1"].sessions) + registry.mu.Unlock() + if active != 2 { + t.Fatalf("expected 2 active sessions, got %d", active) + } + + if _, _, err := registry.ReserveSession("runner-1"); !isRunnerErrorStatus(err, http.StatusServiceUnavailable) { + t.Fatalf("expected no capacity runner error, got %v", err) + } +} + +func TestLiveRunnerRegistry_ReserveSessionUsesProvidedID(t *testing.T) { + registry := newLiveRunnerTestRegistry() + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + + sessionID, endpoint, err := registry.ReserveSession("runner-1", "manifest-id") + if err != nil { + t.Fatal(err) + } + if sessionID != "manifest-id" { + t.Fatalf("expected provided session id, got %q", sessionID) + } + if endpoint != "https://runner.example.com" { + t.Fatalf("unexpected endpoint: %s", endpoint) + } +} + +func TestLiveRunnerRegistry_ReleaseSessionFreesCapacity(t *testing.T) { + registry := newLiveRunnerTestRegistry() + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + if _, _, err := registry.ReserveSession("runner-1"); !isRunnerErrorStatus(err, http.StatusServiceUnavailable) { + t.Fatalf("expected no capacity runner error, got %v", err) + } + if err := registry.ReleaseSession("runner-1", sessionID); err != nil { + t.Fatal(err) + } + if _, _, err := registry.ReserveSession("runner-1"); err != nil { + t.Fatalf("expected capacity after release, got %v", err) + } +} + +func TestLiveRunnerRegistry_SessionToken(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-1") + req.Capacity = 2 + liveRunnerTestRegister(t, registry, req) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + + // Check session token encoding and validation. + token, err := registry.SessionTokenForSession("runner-1", sessionID) + if err != nil { + t.Fatal(err) + } + if !liveRunnerTestSecretPattern.MatchString(token) { + t.Fatalf("expected base32 session token, got %q", token) + } + if err := registry.ValidSessionToken("runner-1", sessionID, token); err != nil { + t.Fatal(err) + } + + otherSessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + // Check token is scoped to its session. + if err := registry.ValidSessionToken("runner-1", otherSessionID, token); !isRunnerErrorStatus(err, http.StatusUnauthorized) { + t.Fatalf("expected token to be scoped to one session, got %v", err) + } +} + +func TestLiveRunnerRegistry_ConcurrentReservationsAreRunnerScoped(t *testing.T) { + registry := newLiveRunnerTestRegistry() + const requestsPerRunner = 64 + for _, runnerID := range []string{"runner-a", "runner-b"} { + req := liveRunnerTestHeartbeat(runnerID) + req.RunnerURL = fmt.Sprintf("https://%s.example.com", runnerID) + req.Capacity = requestsPerRunner + liveRunnerTestRegister(t, registry, req) + } + + registry.mu.Lock() + runnerA := registry.runners["runner-a"] + if runnerA == nil { + registry.mu.Unlock() + t.Fatal("expected runner-a to be registered") + } + runnerA.mu.Lock() + + var runnerAWG sync.WaitGroup + runnerAErrCh := make(chan error, requestsPerRunner) + startRunnerA := make(chan struct{}) + for i := 0; i < requestsPerRunner; i++ { + runnerAWG.Add(1) + go func(i int) { + defer runnerAWG.Done() + <-startRunnerA + sessionID := fmt.Sprintf("runner-a-session-%d", i) + gotSessionID, endpoint, err := registry.ReserveSession("runner-a", sessionID) + if err != nil { + runnerAErrCh <- fmt.Errorf("runner-a reserve %s: %w", sessionID, err) + return + } + if gotSessionID != sessionID { + runnerAErrCh <- fmt.Errorf("runner-a reserve got session %q want %q", gotSessionID, sessionID) + return + } + if endpoint != "https://runner-a.example.com" { + runnerAErrCh <- fmt.Errorf("runner-a endpoint got %q want %q", endpoint, "https://runner-a.example.com") + } + }(i) + } + close(startRunnerA) + registry.mu.Unlock() + + time.Sleep(10 * time.Millisecond) + + var runnerBWG sync.WaitGroup + runnerBErrCh := make(chan error, requestsPerRunner) + runnerBDone := make(chan struct{}) + for i := 0; i < requestsPerRunner; i++ { + runnerBWG.Add(1) + go func(i int) { + defer runnerBWG.Done() + sessionID := fmt.Sprintf("runner-b-session-%d", i) + gotSessionID, endpoint, err := registry.ReserveSession("runner-b", sessionID) + if err != nil { + runnerBErrCh <- fmt.Errorf("runner-b reserve %s: %w", sessionID, err) + return + } + if gotSessionID != sessionID { + runnerBErrCh <- fmt.Errorf("runner-b reserve got session %q want %q", gotSessionID, sessionID) + return + } + if endpoint != "https://runner-b.example.com" { + runnerBErrCh <- fmt.Errorf("runner-b endpoint got %q want %q", endpoint, "https://runner-b.example.com") + } + }(i) + } + go func() { + runnerBWG.Wait() + close(runnerBDone) + }() + + select { + case <-runnerBDone: + case <-time.After(250 * time.Millisecond): + runnerA.mu.Unlock() + t.Fatal("runner-b reservations blocked while runner-a operations were waiting on runner-a.mu") + } + close(runnerBErrCh) + for err := range runnerBErrCh { + if err != nil { + runnerA.mu.Unlock() + t.Fatal(err) + } + } + + runnerADone := make(chan struct{}) + go func() { + runnerAWG.Wait() + close(runnerADone) + }() + runnerA.mu.Unlock() + select { + case <-runnerADone: + case <-time.After(250 * time.Millisecond): + t.Fatal("runner-a reservations did not complete after runner-a.mu was released") + } + close(runnerAErrCh) + for err := range runnerAErrCh { + if err != nil { + t.Fatal(err) + } + } + + runners := registry.Runners() + if len(runners) != 2 { + t.Fatalf("expected two runners, got %+v", runners) + } + usedByURL := make(map[string]int, 2) + availableByURL := make(map[string]int, 2) + for _, runner := range runners { + usedByURL[runner.URL] = runner.CapacityUsed + availableByURL[runner.URL] = runner.CapacityAvailable + } + for _, runnerID := range []string{"runner-a", "runner-b"} { + url := fmt.Sprintf("http://localhost:1234/apps/%s/session", runnerID) + if usedByURL[url] != requestsPerRunner { + t.Fatalf("%s used capacity got %d want %d", runnerID, usedByURL[url], requestsPerRunner) + } + if availableByURL[url] != 0 { + t.Fatalf("%s available capacity got %d want %d", runnerID, availableByURL[url], 0) + } + } +} + +func TestLiveRunnerRegistry_UnregisterRacesWithSessionActions(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-race") + req.Capacity = 100 + resp := liveRunnerTestRegister(t, registry, req) + sessionID, _, err := registry.ReserveSession("runner-race") + if err != nil { + t.Fatal(err) + } + + var wg sync.WaitGroup + errCh := make(chan error, 101) + for i := 0; i < 100; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + switch i % 3 { + case 0: + _, err := registry.RunnerEndpointForSession("runner-race", sessionID) + if err != nil && !isRunnerErrorStatus(err, http.StatusNotFound) { + errCh <- fmt.Errorf("endpoint action: %w", err) + } + case 1: + err := registry.ValidSessionToken("runner-race", sessionID, "bad-token") + if err != nil && !isRunnerErrorStatus(err, http.StatusUnauthorized) && !isRunnerErrorStatus(err, http.StatusNotFound) { + errCh <- fmt.Errorf("token action: %w", err) + } + default: + if err := registry.ReleaseSession("runner-race", sessionID); err != nil && !isRunnerErrorStatus(err, http.StatusNotFound) { + errCh <- fmt.Errorf("release action: %w", err) + } + } + }(i) + } + wg.Add(1) + go func() { + defer wg.Done() + if err := registry.Unregister("runner-race", resp.HeartbeatSecret); err != nil && !isRunnerErrorStatus(err, http.StatusNotFound) { + errCh <- fmt.Errorf("unregister: %w", err) + } + }() + wg.Wait() + close(errCh) + for err := range errCh { + if err != nil { + t.Fatal(err) + } + } + if _, err := registry.RunnerEndpointForSession("runner-race", sessionID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected runner to be gone after unregister, got %v", err) + } +} + +func TestLiveRunnerRegistry_CreateTrickleChannelForSession(t *testing.T) { + trickleSrv, channelBaseURL, _ := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + + channel, err := registry.CreateTrickleChannel("runner-1", sessionID, "foo/bar", "video/MP2T") + if err != nil { + t.Fatal(err) + } + expectedName := sessionID + "-foo-bar" + if channel.ChannelName != expectedName { + t.Fatalf("unexpected channel name channel=%+v want=%s", channel, expectedName) + } + if channel.Name != "foo-bar" { + t.Fatalf("unexpected sanitized name: %s", channel.Name) + } + if channel.URL != channelBaseURL+expectedName { + t.Fatalf("unexpected channel URL: %s", channel.URL) + } +} + +func TestLiveRunnerRegistry_CreateTrickleChannelReturnsExisting(t *testing.T) { + trickleSrv, channelBaseURL, _ := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + + first, err := registry.CreateTrickleChannel("runner-1", sessionID, "existing", "video/MP2T") + if err != nil { + t.Fatal(err) + } + second, err := registry.CreateTrickleChannel("runner-1", sessionID, "existing", "application/json") + if err != nil { + t.Fatal(err) + } + if first.URL != second.URL || second.MimeType != "video/MP2T" { + t.Fatalf("expected existing channel metadata, first=%+v second=%+v", first, second) + } +} + +func TestLiveRunnerRegistry_CreateTrickleChannelRejectsInvalidName(t *testing.T) { + trickleSrv, channelBaseURL, _ := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + + if _, err := registry.CreateTrickleChannel("runner-1", sessionID, "bad.name", ""); !isRunnerErrorStatus(err, http.StatusBadRequest) { + t.Fatalf("expected bad request for invalid channel name, got %v", err) + } + if _, err := registry.CreateTrickleChannel("runner-1", sessionID, "", ""); !isRunnerErrorStatus(err, http.StatusBadRequest) { + t.Fatalf("expected bad request for empty channel name, got %v", err) + } +} + +func TestLiveRunnerRegistry_TrickleChannelCleanup(t *testing.T) { + trickleSrv, channelBaseURL, channelStatus := newLiveRunnerTestTrickleServer(t) + registry := newLiveRunnerTestRegistry() + registry.SetTrickleServer(trickleSrv, channelBaseURL, channelBaseURL) + resp1 := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + resp1Channels := liveRunnerTestBootstrapChannels(resp1) + if len(resp1Channels) != 1 { + t.Fatalf("expected bootstrap trickle channels, got %+v", resp1) + } + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + + if _, err := registry.CreateTrickleChannel("runner-1", sessionID, "cleanup", ""); err != nil { + t.Fatal(err) + } + if status := channelStatus(sessionID + "-cleanup"); status != http.StatusOK { + t.Fatalf("expected channel to exist, got status=%d", status) + } + if err := registry.ReleaseSession("runner-1", sessionID); err != nil { + t.Fatal(err) + } + if status := channelStatus(sessionID + "-cleanup"); status != http.StatusNotFound { + t.Fatalf("expected release to close channel, got status=%d", status) + } + + sessionID, _, err = registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + if _, err := registry.CreateTrickleChannel("runner-1", sessionID, "cleanup", ""); err != nil { + t.Fatal(err) + } + if err := registry.Unregister("runner-1", resp1.HeartbeatSecret); err != nil { + t.Fatal(err) + } + if status := channelStatus(sessionID + "-cleanup"); status != http.StatusNotFound { + t.Fatalf("expected unregister to close channel, got status=%d", status) + } + for _, channel := range resp1Channels { + if status := channelStatus(channel.ChannelName); status != http.StatusNotFound { + t.Fatalf("expected unregister to close bootstrap channel %q, got status=%d", channel.ChannelName, status) + } + } + + registry.heartbeatTTL = defaultLiveRunnerHeartbeatTTL + resp2 := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-2")) + resp2Channels := liveRunnerTestBootstrapChannels(resp2) + sessionID, _, err = registry.ReserveSession("runner-2") + if err != nil { + t.Fatal(err) + } + if _, err := registry.CreateTrickleChannel("runner-2", sessionID, "cleanup", ""); err != nil { + t.Fatal(err) + } + registry.heartbeatTTL = time.Nanosecond + time.Sleep(time.Millisecond) + if _, err := registry.RunnerEndpointForSession("runner-2", sessionID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected expired runner to be unusable, got %v", err) + } + if status := channelStatus(sessionID + "-cleanup"); status != http.StatusOK { + t.Fatalf("expected expired runner cleanup to be asynchronous, got status=%d", status) + } + registry.removeExpiredRunners(time.Now()) + if status := channelStatus(sessionID + "-cleanup"); status != http.StatusNotFound { + t.Fatalf("expected expiry to close channel, got status=%d", status) + } + for _, channel := range resp2Channels { + if status := channelStatus(channel.ChannelName); status != http.StatusNotFound { + t.Fatalf("expected expiry to close bootstrap channel %q, got status=%d", channel.ChannelName, status) + } + } +} + +func TestLiveRunnerRegistry_ExpireAndUnregisterClearSessions(t *testing.T) { + registry := newLiveRunnerTestRegistry() + liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-1")) + sessionID, _, err := registry.ReserveSession("runner-1") + if err != nil { + t.Fatal(err) + } + registry.heartbeatTTL = time.Nanosecond + time.Sleep(time.Millisecond) + if _, err := registry.RunnerEndpointForSession("runner-1", sessionID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected expired runner to be unusable, got %v", err) + } + if runners := registry.Runners(); len(runners) != 0 { + t.Fatalf("expected expired runner to be hidden from discovery, got %d", len(runners)) + } + registry.removeExpiredRunners(time.Now()) + + registry.heartbeatTTL = defaultLiveRunnerHeartbeatTTL + resp2 := liveRunnerTestRegister(t, registry, liveRunnerTestHeartbeat("runner-2")) + if _, _, err := registry.ReserveSession("runner-2"); err != nil { + t.Fatal(err) + } + if err := registry.Unregister("runner-2", resp2.HeartbeatSecret); err != nil { + t.Fatal(err) + } + if _, err := registry.RunnerEndpointForSession("runner-2", sessionID); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected unregistered runner to be removed, got %v", err) + } +} + +func TestLiveRunnerRegistry_LockedRunnerDoesNotBlockUnrelatedRunnerLookup(t *testing.T) { + registry := newLiveRunnerTestRegistry() + defer registry.Stop() + + req1 := liveRunnerTestHeartbeat("runner-blocked") + req1.Capacity = 100 + liveRunnerTestRegister(t, registry, req1) + req2 := liveRunnerTestHeartbeat("runner-free") + req2.Capacity = 100 + liveRunnerTestRegister(t, registry, req2) + + registry.mu.Lock() + blockedRunner := registry.runners["runner-blocked"] + registry.mu.Unlock() + + blockedRunner.mu.Lock() + const blockedLookups = 50 + const freeLookups = 50 + startBlocked := make(chan struct{}) + blockedDone := make(chan error, blockedLookups) + for i := 0; i < blockedLookups; i++ { + go func() { + <-startBlocked + _, _, err := registry.ReserveSession("runner-blocked") + blockedDone <- err + }() + } + close(startBlocked) + time.Sleep(10 * time.Millisecond) + + freeDone := make(chan error, freeLookups) + for i := 0; i < freeLookups; i++ { + go func(i int) { + _, _, err := registry.ReserveSession("runner-free", fmt.Sprintf("free-%d", i)) + freeDone <- err + }(i) + } + + for i := 0; i < freeLookups; i++ { + select { + case err := <-freeDone: + if err != nil { + t.Fatalf("expected unrelated runner lookup to complete: %v", err) + } + case <-time.After(100 * time.Millisecond): + t.Fatal("unrelated runner lookup blocked behind locked runner") + } + } + + blockedRunner.mu.Unlock() + for i := 0; i < blockedLookups; i++ { + if err := <-blockedDone; err != nil && !isRunnerErrorStatus(err, http.StatusServiceUnavailable) { + t.Fatalf("expected blocked runner lookup to complete after unlock: %v", err) + } + } +} + +func TestLiveRunnerRegistry_NotReadyCannotReserveSession(t *testing.T) { + registry := newLiveRunnerTestRegistry() + req := liveRunnerTestHeartbeat("runner-1") + req.Status = "busy" + liveRunnerTestRegister(t, registry, req) + if _, _, err := registry.ReserveSession("runner-1"); !isRunnerErrorStatus(err, http.StatusNotFound) { + t.Fatalf("expected not-ready runner to be unavailable, got %v", err) + } +} + +func isRunnerErrorStatus(err error, statusCode int) bool { + var runnerErr *RunnerError + return errors.As(err, &runnerErr) && runnerErr.StatusCode == statusCode +} + +func newLiveRunnerTestTrickleServer(t *testing.T) (*trickle.Server, string, func(string) int) { + t.Helper() + mux := http.NewServeMux() + trickleSrv := trickle.ConfigureServer(trickle.TrickleServerConfig{ + Mux: mux, + BasePath: "/ai/trickle/", + }) + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + + channelBaseURL := ts.URL + "/ai/trickle/" + channelStatus := func(channelName string) int { + t.Helper() + resp, err := http.Get(channelBaseURL + channelName + "/next") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + return resp.StatusCode + } + return trickleSrv, channelBaseURL, channelStatus +} + +func liveRunnerTestSetO2RKeepaliveInterval(t *testing.T, interval time.Duration) func() { + t.Helper() + previous := liveRunnerO2RKeepaliveInterval + liveRunnerO2RKeepaliveInterval = interval + return func() { + liveRunnerO2RKeepaliveInterval = previous + } +} + +func liveRunnerTestReadTrickleMessage(t *testing.T, trickleSrv *trickle.Server, channelName string, match func(string) bool) string { + t.Helper() + sub := trickle.NewLocalSubscriber(trickleSrv, channelName) + sub.SetSeq(0) + deadline := time.Now().Add(time.Second) + var lastErr error + for time.Now().Before(deadline) { + td, err := sub.Read() + if err != nil { + lastErr = err + time.Sleep(10 * time.Millisecond) + continue + } + data, err := io.ReadAll(td.Reader) + if err != nil { + t.Fatal(err) + } + msg := string(data) + if match(msg) { + return msg + } + } + t.Fatalf("timed out reading matching trickle message from %q, lastErr=%v", channelName, lastErr) + return "" +} + +func liveRunnerTestReadSessionEvent(t *testing.T, trickleSrv *trickle.Server, channelName, event, sessionID string) { + t.Helper() + msg := liveRunnerTestReadTrickleMessage(t, trickleSrv, channelName, func(msg string) bool { + var got struct { + Event string `json:"event"` + Session string `json:"session"` + Timestamp time.Time `json:"timestamp"` + } + if err := json.Unmarshal([]byte(msg), &got); err != nil { + return false + } + return got.Event == event && got.Session == sessionID && !got.Timestamp.IsZero() + }) + var got struct { + Event string `json:"event"` + Session string `json:"session"` + Timestamp time.Time `json:"timestamp"` + } + if err := json.Unmarshal([]byte(msg), &got); err != nil { + t.Fatal(err) + } + if got.Event != event || got.Session != sessionID || got.Timestamp.IsZero() { + t.Fatalf("unexpected session event: %+v", got) + } +} diff --git a/ai/worker/serverless_worker.go b/ai/worker/serverless_worker.go new file mode 100644 index 0000000000..e72178e79f --- /dev/null +++ b/ai/worker/serverless_worker.go @@ -0,0 +1,975 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + "github.com/livepeer/go-livepeer/clog" + "github.com/livepeer/go-livepeer/trickle" +) + +type ServerlessWorker struct { + mu sync.Mutex + inUse int + capacity int + wsURL string + trickleSrv *trickle.Server +} + +type wsMessage struct { + Type string `json:"type"` +} + +type wsHandshakeMessage struct { + Type string `json:"type"` + Code string `json:"code,omitempty"` + Error string `json:"error,omitempty"` + Message string `json:"message,omitempty"` +} + +// ServerlessHandshakeError preserves structured websocket handshake failures. +type ServerlessHandshakeError struct { + StatusCode int + Code string + Message string +} + +func (e *ServerlessHandshakeError) Error() string { + if e == nil { + return "serverless handshake failed" + } + if e.StatusCode == http.StatusBadRequest && e.Message != "" { + return e.Message + } + if e.Code != "" && e.Message != "" { + return fmt.Sprintf("serverless handshake failed (%s): %s", e.Code, e.Message) + } + if e.Message != "" { + return fmt.Sprintf("serverless handshake failed: %s", e.Message) + } + if e.Code != "" { + return fmt.Sprintf("serverless handshake failed (%s)", e.Code) + } + return "serverless handshake failed" +} + +type createChannelsMessage struct { + Type string `json:"type"` + RequestID string `json:"request_id"` + MimeType string `json:"mime_type"` + Direction string `json:"direction"` +} + +type closeChannelsMessage struct { + Type string `json:"type"` + Channels []string `json:"channels"` +} + +type wsResponseMessage struct { + Type string `json:"type"` + RequestID string `json:"request_id"` + Channels []channelInfo `json:"channels"` +} + +type wsPingMessage struct { + Type string `json:"type"` + RequestID string `json:"request_id,omitempty"` + Timestamp int64 `json:"timestamp,omitempty"` +} + +type wsRestartingMessage struct { + Type string `json:"type"` + RequestID string `json:"request_id"` +} + +type wsClosingMessage struct { + Type string `json:"type"` + Message string `json:"message,omitempty"` +} + +type channelInfo struct { + URL string `json:"url"` + Direction string `json:"direction"` + MimeType string `json:"mime_type"` +} + +// NewServerlessWorker creates a new ServerlessWorker instance +func NewServerlessWorker(wsURL string, capacity int) (*ServerlessWorker, error) { + if wsURL == "" { + return nil, fmt.Errorf("serverless worker requires a websocket URL") + } + if capacity <= 0 { + return nil, fmt.Errorf("serverless worker requires a positive capacity") + } + + return &ServerlessWorker{ + capacity: capacity, + wsURL: wsURL, + }, nil +} + +func getServerlessTimeout() time.Duration { + ctx := context.Background() + timeout := 1 * time.Hour + timeoutEnv := strings.TrimSpace(os.Getenv("LIVE_AI_SERVERLESS_TIMEOUT")) + if timeoutEnv == "" { + return timeout + } + + parsedTimeout, err := time.ParseDuration(timeoutEnv) + if err != nil { + clog.Warning(ctx, "Invalid serverless timeout, using default", "timeout", timeoutEnv, "default", timeout, "error", err) + return timeout + } + if parsedTimeout <= 0 { + clog.Warning(ctx, "Serverless timeout must be positive, using default", "timeout", timeoutEnv, "default", timeout) + return timeout + } + + return parsedTimeout +} + +func getServerlessTimeoutGracePeriod() time.Duration { + ctx := context.Background() + gracePeriod := 5 * time.Second + gracePeriodEnv := strings.TrimSpace(os.Getenv("LIVE_AI_SERVERLESS_TIMEOUT_GRACE_PERIOD")) + if gracePeriodEnv == "" { + return gracePeriod + } + + parsedGracePeriod, err := time.ParseDuration(gracePeriodEnv) + if err != nil { + clog.Warning(ctx, "Invalid serverless timeout grace period, using default", "grace_period", gracePeriodEnv, "default", gracePeriod, "error", err) + return gracePeriod + } + if parsedGracePeriod <= 0 { + clog.Warning(ctx, "Serverless timeout grace period must be positive, using default", "grace_period", gracePeriodEnv, "default", gracePeriod) + return gracePeriod + } + + return parsedGracePeriod +} + +func (f *ServerlessWorker) SetTrickleServer(srv *trickle.Server) { + f.mu.Lock() + defer f.mu.Unlock() + f.trickleSrv = srv +} + +func (f *ServerlessWorker) TextToImage(ctx context.Context, req GenTextToImageJSONRequestBody) (*ImageResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.TextToImage not implemented") +} + +func (f *ServerlessWorker) ImageToImage(ctx context.Context, req GenImageToImageMultipartRequestBody) (*ImageResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.ImageToImage not implemented") +} + +func (f *ServerlessWorker) ImageToVideo(ctx context.Context, req GenImageToVideoMultipartRequestBody) (*VideoResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.ImageToVideo not implemented") +} + +func (f *ServerlessWorker) Upscale(ctx context.Context, req GenUpscaleMultipartRequestBody) (*ImageResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.Upscale not implemented") +} + +func (f *ServerlessWorker) AudioToText(ctx context.Context, req GenAudioToTextMultipartRequestBody) (*TextResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.AudioToText not implemented") +} + +func (f *ServerlessWorker) LLM(ctx context.Context, req GenLLMJSONRequestBody) (interface{}, error) { + return nil, fmt.Errorf("ServerlessWorker.LLM not implemented") +} + +func (f *ServerlessWorker) SegmentAnything2(ctx context.Context, req GenSegmentAnything2MultipartRequestBody) (*MasksResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.SegmentAnything2 not implemented") +} + +func (f *ServerlessWorker) ImageToText(ctx context.Context, req GenImageToTextMultipartRequestBody) (*ImageToTextResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.ImageToText not implemented") +} + +func (f *ServerlessWorker) TextToSpeech(ctx context.Context, req GenTextToSpeechJSONRequestBody) (*AudioResponse, error) { + return nil, fmt.Errorf("ServerlessWorker.TextToSpeech not implemented") +} + +func (f *ServerlessWorker) LiveVideoToVideo(ctx context.Context, req GenLiveVideoToVideoJSONRequestBody) (*LiveVideoToVideoResponse, error) { + manifestID := "" + if req.ManifestId != nil { + manifestID = *req.ManifestId + } + if manifestID == "" { + return nil, fmt.Errorf("serverless worker requires manifest ID for live-video-to-video") + } + + if req.ControlUrl == nil || *req.ControlUrl == "" { + return nil, fmt.Errorf("serverless worker requires control URL for live-video-to-video") + } + controlURL, err := url.Parse(*req.ControlUrl) + if err != nil { + return nil, fmt.Errorf("invalid control URL: %w", err) + } + controlPathSuffix := manifestID + "-control" + if !strings.HasSuffix(controlURL.Path, controlPathSuffix) { + return nil, fmt.Errorf("control URL path %q must end with manifest ID suffix %q", controlURL.Path, controlPathSuffix) + } + trickleBasePath := strings.TrimSuffix(controlURL.Path, controlPathSuffix) + buildTrickleURL := func(channelName string) string { + if trickleBasePath == "" { + return "" + } + u := *controlURL + u.Path = trickleBasePath + u = *u.JoinPath(channelName) + return u.String() + } + if buildTrickleURL(manifestID+"-test") == "" { + return nil, fmt.Errorf("could not construct trickle channel URL from control URL %q", *req.ControlUrl) + } + + wsURL := f.wsURL + wsURLPrefix := strings.ToLower(strings.TrimSpace(os.Getenv("LIVE_AI_WS_PREFIX"))) + if req.Params != nil && wsURLPrefix != "" { + if wsURLRaw, ok := (*req.Params)["ws_url"]; ok { + wsURLOverride, ok := wsURLRaw.(string) + if !ok { + return nil, &ServerlessHandshakeError{ + StatusCode: http.StatusBadRequest, + Code: "INVALID_WS", + Message: "invalid params.ws_url: must be a string", + } + } + if wsURLOverride != "" { + wsURLOverride = strings.ToLower(wsURLOverride) + if !strings.HasPrefix(wsURLOverride, wsURLPrefix) { + return nil, &ServerlessHandshakeError{ + StatusCode: http.StatusBadRequest, + Code: "INVALID_WS", + Message: "unacceptable params.ws_url", + } + } + parsedWSURL, err := url.Parse(wsURLOverride) + if err != nil || parsedWSURL.Host == "" { + return nil, &ServerlessHandshakeError{ + StatusCode: http.StatusBadRequest, + Code: "INVALID_WS", + Message: "invalid params.ws_url", + } + } + wsURL = wsURLOverride + } + delete(*req.Params, "ws_url") + } + } + + daydreamUserID := "daydreamer" + if req.Params != nil { + if rawUserID, ok := (*req.Params)["daydream_user_id"]; ok { + userIDStr, ok := rawUserID.(string) + if !ok { + clog.Warning(ctx, "Invalid params.daydream_user_id (not a string), using default", "default", daydreamUserID) + } else if userIDStr == "" { + clog.Warning(ctx, "Empty params.daydream_user_id, using default", "default", daydreamUserID) + } else { + daydreamUserID = userIDStr + } + // NB: the runner checks daydream_user_id so leave it in the params + } + } + + f.mu.Lock() + trickleSrv := f.trickleSrv + if trickleSrv == nil { + f.mu.Unlock() + return nil, fmt.Errorf("serverless worker requires trickle server before live-video-to-video") + } + f.inUse++ + currentInUse := f.inUse + f.mu.Unlock() + + clog.Info(ctx, "Starting new job", "inUse", currentInUse, "capacity", f.capacity, "url", wsURL) + + // Prepare headers with authorization + headers := http.Header{} + if authToken := os.Getenv("FAL_API_KEY"); authToken != "" { + headers.Add("Authorization", "Key "+authToken) + clog.Info(ctx, "Added authorization header from FAL_API_KEY") + } + headers.Set("Manifest-Id", manifestID) + headers.Set("Daydream-User-Id", daydreamUserID) + + // Dial before starting goroutines so errors can be returned to caller + rawWebsocketConn, _, err := websocket.DefaultDialer.Dial(wsURL, headers) + if err != nil { + f.mu.Lock() + f.inUse-- + remainingInUse := f.inUse + f.mu.Unlock() + clog.Error(ctx, "Failed to connect to websocket", "error", err, "inUse", remainingInUse, "capacity", f.capacity) + return nil, fmt.Errorf("failed to connect to websocket: %w", err) + } + + // Marshal request to JSON and send. + reqJSON, err := json.Marshal(req) + if err != nil { + _ = rawWebsocketConn.Close() + f.mu.Lock() + f.inUse-- + remainingInUse := f.inUse + f.mu.Unlock() + clog.Error(ctx, "Failed to marshal request", "error", err, "inUse", remainingInUse, "capacity", f.capacity) + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + websocketConn := newLockedWebSocket(rawWebsocketConn) + + // Wait for connection readiness before starting background processing or returning. + _, readyMsg, err := websocketConn.ReadMessage() + if err != nil { + _ = websocketConn.Close() + f.mu.Lock() + f.inUse-- + remainingInUse := f.inUse + f.mu.Unlock() + clog.Error(ctx, "Failed to read ready message", "error", err, "inUse", remainingInUse, "capacity", f.capacity) + return nil, fmt.Errorf("failed to read ready message: %w", err) + } + + err = performHandshake( + ctx, + readyMsg, + func() ([]byte, error) { + _, msg, readErr := websocketConn.ReadMessage() + return msg, readErr + }, + websocketConn, + reqJSON, + ) + if err != nil { + _ = websocketConn.Close() + f.mu.Lock() + f.inUse-- + remainingInUse := f.inUse + f.mu.Unlock() + clog.Error(ctx, "Failed handshake", "error", err, "message", string(readyMsg), "inUse", remainingInUse, "capacity", f.capacity) + return nil, err + } + + // Start websocket processing in a goroutine + go func() { + openChannels := map[string]*trickle.TrickleLocalPublisher{} + + // Ensure we decrement the counter and close all channels when this goroutine exits + defer func() { + _ = websocketConn.Close() + + for name, ch := range openChannels { + if err := ch.Close(); err != nil { + clog.Warning(ctx, "Failed to close trickle channel", "channel", name, "error", err) + } + } + eventsChannelName := manifestID + "-events" + eventsCh := trickle.NewLocalPublisher(trickleSrv, eventsChannelName, "application/json") + if err := eventsCh.Close(); err != nil { + clog.Warning(ctx, "Failed to close events trickle channel", "channel", eventsChannelName, "error", err) + } + controlChannelName := manifestID + "-control" + controlCh := trickle.NewLocalPublisher(trickleSrv, controlChannelName, "application/json") + if err := controlCh.Close(); err != nil { + clog.Warning(ctx, "Failed to close control trickle channel", "channel", controlChannelName, "error", err) + } + + f.mu.Lock() + f.inUse-- + remainingInUse := f.inUse + f.mu.Unlock() + clog.Info(ctx, "Job ended", "inUse", remainingInUse, "capacity", f.capacity) + }() + + clog.Info(ctx, "Connected to websocket successfully") + + // Subscribe to the events trickle stream to detect when stream ends + if req.EventsUrl != nil && *req.EventsUrl != "" { + go func() { + eventsUrl := *req.EventsUrl + clog.Info(ctx, "Subscribing to events stream", "url", eventsUrl) + + subscriber, err := trickle.NewTrickleSubscriber(trickle.TrickleSubscriberConfig{ + URL: eventsUrl, + }) + if err != nil { + clog.Error(ctx, "Failed to create events subscriber", "error", err) + return + } + + const maxRetries = 5 + const retryPause = 300 * time.Millisecond + retries := 0 + + for { + segment, err := subscriber.Read() + if err == nil { + retries = 0 + } else { + if errors.Is(err, trickle.EOS) || errors.Is(err, trickle.StreamNotFoundErr) { + clog.Info(ctx, "Events stream closed, closing websocket", "reason", err) + _ = websocketConn.Close() // This will cause ReadMessage to return an error + return + } + + var seqErr *trickle.SequenceNonexistent + if errors.As(err, &seqErr) { + // Stream exists but segment doesn't, so skip to leading edge. + subscriber.SetSeq(seqErr.Latest) + } + + if retries > maxRetries { + clog.Error(ctx, "Too many errors reading events stream, closing websocket", "error", err) + _ = websocketConn.Close() + return + } + + clog.Warning(ctx, "Error reading events stream", "error", err, "retry", retries) + retries++ + time.Sleep(retryPause) + continue + } + + // Read the event data to drain the stream, but avoid logging full payload contents. + data, err := io.ReadAll(segment.Body) + _ = segment.Body.Close() + if err != nil { + clog.Warning(ctx, "Error reading event body", "error", err) + continue + } + clog.Info(ctx, "Received event from trickle stream", "size_bytes", len(data)) + } + }() + } else { + clog.Warning(ctx, "No events URL provided, cannot detect stream end via trickle") + } + + // Track missed pongs and terminate after 3 missed pongs. + const ( + pingInterval = 10 * time.Second + maxMissedPongs = int32(3) + ) + var missedPongs atomic.Int32 + stopPings := startPinging( + ctx, + pingInterval, + maxMissedPongs, + &missedPongs, + websocketConn, + ) + defer func() { + if stopPings != nil { + stopPings() + } + }() + + // Fail-safe timeout to avoid big bills. + serverlessTimeout := getServerlessTimeout() + timeoutGracePeriod := getServerlessTimeoutGracePeriod() + timeout := time.NewTimer(serverlessTimeout) + defer timeout.Stop() + + // Create a channel to receive messages + messageChan := make(chan struct { + messageType int + message []byte + err error + }) + + // Start goroutine to read websocket messages + go func() { + for { + messageType, message, err := websocketConn.ReadMessage() + messageChan <- struct { + messageType int + message []byte + err error + }{messageType, message, err} + if err != nil { + return + } + } + }() + + streamCount := 0 + const ( + // Retries may be expected due to plugin install/uninstall, so be generous. + maxHandshakeRetries = 20 + handshakeRetryWindow = 1 * time.Minute + ) + handshakeRetryTimes := make([]time.Time, 0, maxHandshakeRetries) + + // Receive and handle messages indefinitely (or until events stream closes or timeout) + for { + select { + case <-timeout.C: + clog.Info(ctx, "Websocket connection timeout reached, notifying runner before cleanup", "timeout", serverlessTimeout, "grace_period", timeoutGracePeriod) + timeoutPayload, err := json.Marshal(wsClosingMessage{ + Type: "closing", + Message: "scope session timeout reached; terminating shortly", + }) + if err != nil { + clog.Warning(ctx, "Failed to marshal websocket timeout message before cleanup", "error", err) + return + } + if err := websocketConn.Write(timeoutPayload); err != nil { + clog.Warning(ctx, "Failed to notify websocket timeout before cleanup", "error", err) + return + } + // This blocks the loop, so no more processing happens during the grace period. + time.Sleep(timeoutGracePeriod) + return + + case msg := <-messageChan: + if msg.err != nil { + clog.Info(ctx, "Websocket read ended", "error", msg.err) + return + } + + clog.Info(ctx, "Received message from websocket", + "type", msg.messageType, + "message", string(msg.message)) + + var generic wsMessage + if err := json.Unmarshal(msg.message, &generic); err != nil { + clog.Warning(ctx, "Failed to parse message", "error", err) + continue + } + + switch generic.Type { + case "pong": + missedPongs.Store(0) + + case "ready": + // This is a retry, maybe due to plugin install / uninstall. + + // First check whether this is too many retries in a short time + now := time.Now() + cutoff := now.Add(-handshakeRetryWindow) + filteredRetryTimes := handshakeRetryTimes[:0] + for _, retryAt := range handshakeRetryTimes { + if retryAt.After(cutoff) { + filteredRetryTimes = append(filteredRetryTimes, retryAt) + } + } + handshakeRetryTimes = filteredRetryTimes + if len(handshakeRetryTimes) >= maxHandshakeRetries { + clog.Error(ctx, "Too many handshake retries in short period", "retries", len(handshakeRetryTimes), "window", handshakeRetryWindow) + return + } + handshakeRetryTimes = append(handshakeRetryTimes, now) + + // Close all media channels; active streams do not survive a retry. + for channelName, ch := range openChannels { + if err := ch.Close(); err != nil { + clog.Warning(ctx, "Failed to close media trickle channel before handshake retry", "channel", channelName, "error", err) + continue + } + delete(openChannels, channelName) + } + streamCount = 0 + + // Perform the handshake + err := performHandshake( + ctx, + msg.message, + func() ([]byte, error) { + nextMsg := <-messageChan + if nextMsg.err != nil { + return nil, nextMsg.err + } + return nextMsg.message, nil + }, + websocketConn, + reqJSON, + ) + if err != nil { + clog.Error(ctx, "Handshake retry failed", "error", err) + return + } + if stopPings != nil { + stopPings() + } + stopPings = startPinging( + ctx, + pingInterval, + maxMissedPongs, + &missedPongs, + websocketConn, + ) + missedPongs.Store(0) + clog.Info(ctx, "Handshake retry completed") + + case "restarting": + if stopPings != nil { + stopPings() + stopPings = nil + } + // Hack: Run this in a goroutine so we can accommodate other in-flight + // messages, mostly pongs + go func(rawMessage []byte) { + respJSON, err := handleRestartingMessage(ctx, rawMessage, &missedPongs) + if err != nil { + clog.Warning(ctx, "Failed to handle restarting message", "error", err) + return + } + err = websocketConn.Write(respJSON) + if err != nil { + clog.Warning(ctx, "Failed to send restarting response", "error", err) + } + }(msg.message) + + case "create_channels": + var startMsg createChannelsMessage + if err := json.Unmarshal(msg.message, &startMsg); err != nil { + clog.Warning(ctx, "Failed to parse create_channels message", "error", err) + continue + } + mimeType := startMsg.MimeType + if mimeType == "" { + mimeType = "video/MP2T" + } + + streamCount++ + streamCountStr := strconv.Itoa(streamCount) + var newChannels []channelInfo + createInbound := false + createOutbound := false + switch strings.ToLower(startMsg.Direction) { + case "in": + createInbound = true + case "out": + createOutbound = true + case "bidirectional": + createInbound = true + createOutbound = true + default: + clog.Warning(ctx, "Ignoring create_channels with unsupported direction", "direction", startMsg.Direction) + continue + } + + if createInbound { + channelName := manifestID + "-" + streamCountStr + "-in" + channelURL := buildTrickleURL(channelName) + ch := trickle.NewLocalPublisher(trickleSrv, channelName, mimeType) + ch.CreateChannel() + openChannels[channelURL] = ch + newChannels = append(newChannels, channelInfo{ + URL: channelURL, + Direction: "in", + MimeType: mimeType, + }) + } + if createOutbound { + channelName := manifestID + "-" + streamCountStr + "-out" + channelURL := buildTrickleURL(channelName) + ch := trickle.NewLocalPublisher(trickleSrv, channelName, mimeType) + ch.CreateChannel() + openChannels[channelURL] = ch + newChannels = append(newChannels, channelInfo{ + URL: channelURL, + Direction: "out", + MimeType: mimeType, + }) + } + + resp := wsResponseMessage{ + Type: "response", + RequestID: startMsg.RequestID, + Channels: newChannels, + } + respJSON, err := json.Marshal(resp) + if err != nil { + clog.Warning(ctx, "Failed to marshal response message", "error", err) + continue + } + err = websocketConn.Write(respJSON) + if err != nil { + clog.Warning(ctx, "Failed to send response message", "error", err) + return + } + clog.Info(ctx, "Created trickle channels for create_channels", "count", len(newChannels), "direction", startMsg.Direction, "mimeType", mimeType) + + case "close_channels": + var stopMsg closeChannelsMessage + if err := json.Unmarshal(msg.message, &stopMsg); err != nil { + clog.Warning(ctx, "Failed to parse close_channels message", "error", err) + continue + } + for _, channelName := range stopMsg.Channels { + ch, ok := openChannels[channelName] + if !ok { + clog.Warning(ctx, "close_channels channel not found", "channel", channelName) + continue + } + if err := ch.Close(); err != nil { + clog.Warning(ctx, "Failed to close close_channels channel", "channel", channelName, "error", err) + continue + } + delete(openChannels, channelName) + clog.Info(ctx, "Closed trickle channel from close_channels", "channel", channelName) + } + default: + clog.Error(ctx, "Received unrecognized websocket message type", "type", generic.Type, "message", string(msg.message)) + return + } + } + } + }() + + return &LiveVideoToVideoResponse{}, nil +} + +func (f *ServerlessWorker) Warm(ctx context.Context, pipeline string, modelID string, endpoint RunnerEndpoint, optimizationFlags OptimizationFlags) error { + // Serverless workers don't need warming + return nil +} + +func (f *ServerlessWorker) Stop(ctx context.Context) error { + // No containers to stop for serverless workers + return nil +} + +func (f *ServerlessWorker) HasCapacity(pipeline string, modelID string) bool { + f.mu.Lock() + defer f.mu.Unlock() + + hasCapacity := f.inUse < f.capacity + clog.Info(context.Background(), "HasCapacity", + "pipeline", pipeline, + "modelID", modelID, + "inUse", f.inUse, + "capacity", f.capacity, + "hasCapacity", hasCapacity) + + return hasCapacity +} + +func (f *ServerlessWorker) EnsureImageAvailable(ctx context.Context, pipeline string, modelID string) error { + // No images to download for serverless workers + return nil +} + +func (f *ServerlessWorker) HardwareInformation() []HardwareInformation { + // Return empty for serverless workers + return []HardwareInformation{} +} + +func (f *ServerlessWorker) GetLiveAICapacity(pipeline, modelID string) Capacity { + f.mu.Lock() + defer f.mu.Unlock() + + idleCapacity := f.capacity - f.inUse + if idleCapacity < 0 { + idleCapacity = 0 + } + + clog.Info(context.Background(), "GetLiveAICapacity", + "pipeline", pipeline, + "modelID", modelID, + "inUse", f.inUse, + "idleCapacity", idleCapacity) + + return Capacity{ + ContainersInUse: f.inUse, + ContainersIdle: idleCapacity, + } +} + +func (f *ServerlessWorker) Version() []Version { + // Return a default version for serverless workers + return []Version{ + { + Version: "serverless-1.0.0", + }, + } +} + +func performHandshake( + ctx context.Context, + readyMsg []byte, + readMsg func() ([]byte, error), + websocketConn *lockedWebSocket, + reqJSON []byte, +) error { + var readyResponse wsHandshakeMessage + if err := json.Unmarshal(readyMsg, &readyResponse); err != nil { + return fmt.Errorf("failed to parse ready message: %w", err) + } + if readyResponse.Type != "ready" { + return fmt.Errorf("did not receive ready message from websocket") + } + clog.Info(ctx, "Received ready message", "message", readyResponse.Message) + + if len(reqJSON) == 0 { + return fmt.Errorf("empty request payload") + } + clog.Info(ctx, "Sending handshake request", "request", string(reqJSON)) + if err := websocketConn.Write(reqJSON); err != nil { + return fmt.Errorf("failed to send request message: %w", err) + } + + startedMsg, err := readMsg() + if err != nil { + return fmt.Errorf("failed to read started message: %w", err) + } + + var startedResponse wsHandshakeMessage + if err := json.Unmarshal(startedMsg, &startedResponse); err != nil { + return fmt.Errorf("failed to parse started message: %w", err) + } + if startedResponse.Type == "error" { + statusCode := http.StatusInternalServerError + if startedResponse.Code == "ACCESS_DENIED" { + statusCode = http.StatusUnauthorized + } + return &ServerlessHandshakeError{ + StatusCode: statusCode, + Code: startedResponse.Code, + Message: firstNonEmpty(startedResponse.Error, startedResponse.Message), + } + } + if startedResponse.Type != "started" { + return fmt.Errorf("unexpected message type %q between ready and started", startedResponse.Type) + } + clog.Info(ctx, "Received started message") + + return nil +} + +func startPinging( + logCtx context.Context, + pingInterval time.Duration, + maxMissedPongs int32, + missedPongs *atomic.Int32, + websocketConn *lockedWebSocket, +) func() { + ctx, cancel := context.WithCancel(context.Background()) + + go func() { + ticker := time.NewTicker(pingInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + pingPayload := wsPingMessage{ + Type: "ping", + RequestID: fmt.Sprintf("ping-%d", time.Now().UnixNano()), + Timestamp: time.Now().UnixMilli(), + } + pingJSON, err := json.Marshal(pingPayload) + if err != nil { + clog.Warning(logCtx, "Failed to marshal app ping", "error", err) + continue + } + if err := websocketConn.Write(pingJSON); err != nil { + clog.Warning(logCtx, "Failed to send ping", "error", err) + return + } + if missedPongs.Add(1) >= maxMissedPongs { + clog.Warning(logCtx, "Too many missed pongs, closing websocket", "missed", missedPongs.Load()) + _ = websocketConn.Close() + return + } + case <-ctx.Done(): + return + } + } + }() + + return cancel +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func handleRestartingMessage( + ctx context.Context, + message []byte, + missedPongs *atomic.Int32, +) ([]byte, error) { + var restartingMsg wsRestartingMessage + if err := json.Unmarshal(message, &restartingMsg); err != nil { + return nil, fmt.Errorf("failed to parse restarting message: %w", err) + } + + // First wait for any pending pongs so they don't gum up the re-handshake + if currentMissed := missedPongs.Load(); currentMissed > 0 { + const ( + pongWaitAttempts = 3 + pongWaitInterval = 100 * time.Millisecond + ) + for i := 0; i < pongWaitAttempts; i++ { + if missedPongs.Load() == 0 { + clog.Info(ctx, "Recovered from a pending pong, good job") + break + } + time.Sleep(pongWaitInterval) + } + if remainingMissed := missedPongs.Load(); remainingMissed > 0 { + clog.Warning(ctx, "Continuing restart without pong recovery", "missedPongs", remainingMissed) + } + } + missedPongs.Store(0) + + resp := wsResponseMessage{ + Type: "response", + RequestID: restartingMsg.RequestID, + } + respJSON, err := json.Marshal(resp) + if err != nil { + return nil, fmt.Errorf("failed to marshal restarting response: %w", err) + } + + return respJSON, nil +} + +type lockedWebSocket struct { + conn *websocket.Conn + writeMu sync.Mutex + closeOnce sync.Once +} + +func newLockedWebSocket(conn *websocket.Conn) *lockedWebSocket { + return &lockedWebSocket{conn: conn} +} + +func (ws *lockedWebSocket) ReadMessage() (int, []byte, error) { + return ws.conn.ReadMessage() +} + +func (ws *lockedWebSocket) Write(payload []byte) error { + ws.writeMu.Lock() + defer ws.writeMu.Unlock() + return ws.conn.WriteMessage(websocket.TextMessage, payload) +} + +func (ws *lockedWebSocket) Close() error { + var closeErr error + ws.closeOnce.Do(func() { + closeErr = ws.conn.Close() + }) + return closeErr +} diff --git a/ai/worker/serverless_worker_test.go b/ai/worker/serverless_worker_test.go new file mode 100644 index 0000000000..ac5002bebd --- /dev/null +++ b/ai/worker/serverless_worker_test.go @@ -0,0 +1,414 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/gorilla/websocket" +) + +func newTestLockedWebSocket(t *testing.T) (*lockedWebSocket, <-chan []byte, func()) { + t.Helper() + + recv := make(chan []byte, 32) + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + go func() { + defer close(recv) + defer conn.Close() + for { + _, msg, err := conn.ReadMessage() + if err != nil { + return + } + recv <- append([]byte(nil), msg...) + } + }() + })) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + clientConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + server.Close() + t.Fatalf("failed to dial test websocket: %v", err) + } + client := newLockedWebSocket(clientConn) + + cleanup := func() { + _ = client.Close() + server.Close() + } + return client, recv, cleanup +} + +func TestPerformHandshakeSuccess(t *testing.T) { + t.Parallel() + + ws, recv, cleanup := newTestLockedWebSocket(t) + defer cleanup() + + reqJSON := []byte(`{"type":"request"}`) + err := performHandshake( + context.Background(), + []byte(`{"type":"ready","message":"ok"}`), + func() ([]byte, error) { + return []byte(`{"type":"started"}`), nil + }, + ws, + reqJSON, + ) + if err != nil { + t.Fatalf("performHandshake() returned error: %v", err) + } + select { + case sent := <-recv: + if string(sent) != string(reqJSON) { + t.Fatalf("performHandshake() sent wrong request: got=%s want=%s", string(sent), string(reqJSON)) + } + case <-time.After(1 * time.Second): + t.Fatal("performHandshake() did not send request") + } +} + +func TestPerformHandshakeAccessDenied(t *testing.T) { + t.Parallel() + + ws, _, cleanup := newTestLockedWebSocket(t) + defer cleanup() + + err := performHandshake( + context.Background(), + []byte(`{"type":"ready"}`), + func() ([]byte, error) { + return []byte(`{"type":"error","code":"ACCESS_DENIED","error":"Access denied"}`), nil + }, + ws, + []byte(`{"type":"request"}`), + ) + if err == nil { + t.Fatal("performHandshake() expected error, got nil") + } + + var hsErr *ServerlessHandshakeError + if !errors.As(err, &hsErr) { + t.Fatalf("performHandshake() expected ServerlessHandshakeError, got %T: %v", err, err) + } + if hsErr.StatusCode != http.StatusUnauthorized { + t.Fatalf("unexpected status code: got=%d want=%d", hsErr.StatusCode, http.StatusUnauthorized) + } + if hsErr.Code != "ACCESS_DENIED" { + t.Fatalf("unexpected handshake code: %q", hsErr.Code) + } + if hsErr.Message != "Access denied" { + t.Fatalf("unexpected handshake message: %q", hsErr.Message) + } +} + +func TestLiveVideoToVideoRequiresControlURL(t *testing.T) { + t.Parallel() + + worker := &ServerlessWorker{} + manifestID := "manifest" + + _, err := worker.LiveVideoToVideo(t.Context(), GenLiveVideoToVideoJSONRequestBody{ + ManifestId: &manifestID, + SubscribeUrl: "https://example.com/ai/trickle/manifest", + }) + if err == nil { + t.Fatal("LiveVideoToVideo() expected error, got nil") + } + if !strings.Contains(err.Error(), "requires control URL") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestLiveVideoToVideoUsesControlURLForTrickleBase(t *testing.T) { + t.Parallel() + + worker := &ServerlessWorker{} + manifestID := "manifest" + controlURL := "https://example.com/ai/trickle/manifest-control" + + _, err := worker.LiveVideoToVideo(t.Context(), GenLiveVideoToVideoJSONRequestBody{ + ManifestId: &manifestID, + ControlUrl: &controlURL, + SubscribeUrl: "not a valid url", + }) + if err == nil { + t.Fatal("LiveVideoToVideo() expected error, got nil") + } + if !strings.Contains(err.Error(), "requires trickle server") { + t.Fatalf("expected trickle server error after control URL validation, got: %v", err) + } +} + +func TestLiveVideoToVideoInvalidWSURLReturnsBadRequest(t *testing.T) { + tests := []struct { + name string + wsURLValue any + wantErr string + }{ + { + name: "non-string override", + wsURLValue: 123, + wantErr: "invalid params.ws_url: must be a string", + }, + { + name: "unacceptable override", + wsURLValue: "ws://other-host/socket", + wantErr: "unacceptable params.ws_url", + }, + { + name: "invalid override", + wsURLValue: "ws://allowed-host:badport/socket", + wantErr: "invalid params.ws_url", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("LIVE_AI_WS_PREFIX", "ws://allowed-host") + + worker := &ServerlessWorker{wsURL: "ws://default-host/socket"} + manifestID := "manifest" + controlURL := "https://example.com/ai/trickle/manifest-control" + params := map[string]interface{}{"ws_url": tt.wsURLValue} + + _, err := worker.LiveVideoToVideo(t.Context(), GenLiveVideoToVideoJSONRequestBody{ + ManifestId: &manifestID, + ControlUrl: &controlURL, + Params: ¶ms, + }) + if err == nil { + t.Fatal("LiveVideoToVideo() expected error, got nil") + } + + var hsErr *ServerlessHandshakeError + if !errors.As(err, &hsErr) { + t.Fatalf("LiveVideoToVideo() expected ServerlessHandshakeError, got %T: %v", err, err) + } + if hsErr.StatusCode != http.StatusBadRequest { + t.Fatalf("unexpected status code: got=%d want=%d", hsErr.StatusCode, http.StatusBadRequest) + } + if hsErr.Code != "INVALID_WS" { + t.Fatalf("unexpected handshake code: got=%q want=%q", hsErr.Code, "INVALID_WS") + } + if hsErr.Message != tt.wantErr { + t.Fatalf("unexpected handshake message: got=%q want=%q", hsErr.Message, tt.wantErr) + } + if err.Error() != tt.wantErr { + t.Fatalf("unexpected error string: got=%q want=%q", err.Error(), tt.wantErr) + } + if _, ok := params["ws_url"]; !ok { + t.Fatalf("ws_url should remain in params on validation failure: params=%v", params) + } + }) + } +} + +func TestGetServerlessTimeoutDefault(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT", "") + + if got := getServerlessTimeout(); got != time.Hour { + t.Fatalf("unexpected default serverless timeout: got=%v want=%v", got, time.Hour) + } +} + +func TestGetServerlessTimeoutFromEnv(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT", "90m") + + if got := getServerlessTimeout(); got != 90*time.Minute { + t.Fatalf("unexpected configured serverless timeout: got=%v want=%v", got, 90*time.Minute) + } +} + +func TestGetServerlessTimeoutInvalidEnvFallsBack(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT", "not-a-duration") + + if got := getServerlessTimeout(); got != time.Hour { + t.Fatalf("unexpected fallback serverless timeout: got=%v want=%v", got, time.Hour) + } +} + +func TestGetServerlessTimeoutNonPositiveFallsBack(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT", "0s") + + if got := getServerlessTimeout(); got != time.Hour { + t.Fatalf("unexpected fallback serverless timeout for non-positive value: got=%v want=%v", got, time.Hour) + } +} + +func TestGetServerlessTimeoutGracePeriodDefault(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT_GRACE_PERIOD", "") + + if got := getServerlessTimeoutGracePeriod(); got != 5*time.Second { + t.Fatalf("unexpected default serverless timeout grace period: got=%v want=%v", got, 5*time.Second) + } +} + +func TestGetServerlessTimeoutGracePeriodFromEnv(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT_GRACE_PERIOD", "2500ms") + + if got := getServerlessTimeoutGracePeriod(); got != 2500*time.Millisecond { + t.Fatalf("unexpected configured serverless timeout grace period: got=%v want=%v", got, 2500*time.Millisecond) + } +} + +func TestGetServerlessTimeoutGracePeriodInvalidEnvFallsBack(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT_GRACE_PERIOD", "not-a-duration") + + if got := getServerlessTimeoutGracePeriod(); got != 5*time.Second { + t.Fatalf("unexpected fallback serverless timeout grace period: got=%v want=%v", got, 5*time.Second) + } +} + +func TestGetServerlessTimeoutGracePeriodNonPositiveFallsBack(t *testing.T) { + t.Setenv("LIVE_AI_SERVERLESS_TIMEOUT_GRACE_PERIOD", "0s") + + if got := getServerlessTimeoutGracePeriod(); got != 5*time.Second { + t.Fatalf("unexpected fallback serverless timeout grace period for non-positive value: got=%v want=%v", got, 5*time.Second) + } +} + +func TestServerlessClosingMessagePayload(t *testing.T) { + payload, err := json.Marshal(wsClosingMessage{ + Type: "closing", + Message: "live-video-to-video session timeout reached; terminating shortly", + }) + if err != nil { + t.Fatalf("failed to marshal closing message JSON: %v", err) + } + + var msg wsClosingMessage + if err := json.Unmarshal(payload, &msg); err != nil { + t.Fatalf("failed to unmarshal closing message JSON: %v", err) + } + if msg.Type != "closing" { + t.Fatalf("unexpected closing message type: got=%q want=%q", msg.Type, "closing") + } + if msg.Message != "live-video-to-video session timeout reached; terminating shortly" { + t.Fatalf("unexpected closing message text: got=%q", msg.Message) + } +} + +func TestRestartMessage(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + var missedPongs atomic.Int32 + missedPongs.Store(2) + + go func() { + time.Sleep(20 * time.Millisecond) + missedPongs.Store(0) + }() + + start := time.Now() + respJSON, err := handleRestartingMessage( + context.Background(), + []byte(`{"type":"restarting","request_id":"restart-1"}`), + &missedPongs, + ) + elapsed := time.Since(start) + if err != nil { + t.Fatalf("handleRestartingMessage() returned error: %v", err) + } + if elapsed < 100*time.Millisecond { + t.Fatalf("handleRestartingMessage() should wait for pong recovery, elapsed=%v", elapsed) + } + if elapsed >= 200*time.Millisecond { + t.Fatalf("handleRestartingMessage() did not observe concurrent pong recovery, elapsed=%v", elapsed) + } + if got := missedPongs.Load(); got != 0 { + t.Fatalf("handleRestartingMessage() should reset missed pongs, got %d", got) + } + + var resp wsResponseMessage + if err := json.Unmarshal(respJSON, &resp); err != nil { + t.Fatalf("failed to unmarshal restarting response JSON: %v", err) + } + if resp.Type != "response" { + t.Fatalf("unexpected response type: got=%q want=%q", resp.Type, "response") + } + if resp.RequestID != "restart-1" { + t.Fatalf("unexpected request_id: got=%q want=%q", resp.RequestID, "restart-1") + } + }) +} + +func TestStartPingingStopIsIdempotent(t *testing.T) { + ws, _, cleanup := newTestLockedWebSocket(t) + defer cleanup() + + var missedPongs atomic.Int32 + stopPings := startPinging( + context.Background(), + 20*time.Millisecond, + 3, + &missedPongs, + ws, + ) + + stopPings() + stopPings() +} + +func TestStartPingingRestartsWithFullIntervalGrace(t *testing.T) { + ws, _, cleanup := newTestLockedWebSocket(t) + defer cleanup() + + const interval = 40 * time.Millisecond + var missedPongs atomic.Int32 + + stopPings := startPinging( + context.Background(), + interval, + 100, // keep this test focused on timer behavior + &missedPongs, + ws, + ) + + time.Sleep(interval - 10*time.Millisecond) + if got := missedPongs.Load(); got != 0 { + t.Fatalf("expected no ping before first interval, got=%d", got) + } + + // Simulate restart: stop the existing pinger before its next tick. + stopPings() + + // Simulate successful re-handshake: start a fresh pinger. + stopPings = startPinging( + context.Background(), + interval, + 100, + &missedPongs, + ws, + ) + defer stopPings() + + time.Sleep(interval - 10*time.Millisecond) + if got := missedPongs.Load(); got != 0 { + t.Fatalf("expected full interval grace after restart, got=%d", got) + } + + deadline := time.Now().Add(250 * time.Millisecond) + for time.Now().Before(deadline) { + if missedPongs.Load() == 1 { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("expected first ping after full interval, got=%d", missedPongs.Load()) +} diff --git a/byoc/job_orchestrator.go b/byoc/job_orchestrator.go index 210474992e..5a7b28812d 100644 --- a/byoc/job_orchestrator.go +++ b/byoc/job_orchestrator.go @@ -566,18 +566,26 @@ func (bso *BYOCOrchestratorServer) verifyJobCreds(ctx context.Context, jobCreds return nil, errSegSig } - if !bso.orch.VerifySig(ethcommon.HexToAddress(jobData.Sender), jobData.Request+jobData.Parameters, sigByte) { - clog.Errorf(ctx, "Sig check failed sender=%v", jobData.Sender) - return nil, errSegSig - } - - if reserveCapacity && bso.orch.ReserveExternalCapabilityCapacity(jobData.Capability) != nil { - return nil, errZeroCapacity + sender := ethcommon.HexToAddress(jobData.Sender) + + // Verify V1 structured binary format (matches DMZ /sign-byoc-job signing). + v1Payload := FlattenBYOCJob(&BYOCJobSigningInput{ + ID: jobData.ID, + Capability: jobData.Capability, + Request: jobData.Request, + Parameters: jobData.Parameters, + TimeoutSeconds: jobData.Timeout, + }) + if bso.orch.VerifySig(sender, string(v1Payload), sigByte) { + if reserveCapacity && bso.orch.ReserveExternalCapabilityCapacity(jobData.Capability) != nil { + return nil, errZeroCapacity + } + jobData.CapabilityUrl = bso.orch.GetUrlForCapability(jobData.Capability) + return jobData, nil } - jobData.CapabilityUrl = bso.orch.GetUrlForCapability(jobData.Capability) - - return jobData, nil + clog.Errorf(ctx, "Sig check failed sender=%v", jobData.Sender) + return nil, errSegSig } func (bso *BYOCOrchestratorServer) verifyTokenCreds(ctx context.Context, tokenCreds string) (*JobSender, error) { diff --git a/byoc/types.go b/byoc/types.go index fa9b4a5e8b..124a83cd5b 100644 --- a/byoc/types.go +++ b/byoc/types.go @@ -3,6 +3,7 @@ package byoc import ( "context" "crypto/tls" + "encoding/binary" "errors" "math/big" gonet "net" @@ -279,3 +280,65 @@ type byocLiveRequestParams struct { // when the write for the last segment started lastSegmentTime time.Time } + +// Prevents cross-protocol signature replay. +const BYOCJobSigV1Prefix = "LP_BYOC_JOB_V1\x00\x00" + +// BYOCJobSigningInput holds the fields that are bound into a BYOC job signature. +type BYOCJobSigningInput struct { + ID string + Capability string + Request string + Parameters string + TimeoutSeconds int +} + +// FlattenBYOCJob produces a deterministic binary representation of a BYOC job +// for signing, similar to SegTranscodingMetadata.Flatten() used by LV2V. +// +// Wire format: +// +// version(16) || timeout(4,BE) || len(id)(4,BE) || id || len(cap)(4,BE) || cap +// || len(req)(4,BE) || req || len(params)(4,BE) || params +func FlattenBYOCJob(job *BYOCJobSigningInput) []byte { + idBytes := []byte(job.ID) + capBytes := []byte(job.Capability) + reqBytes := []byte(job.Request) + paramsBytes := []byte(job.Parameters) + + size := 16 + 4 + + 4 + len(idBytes) + + 4 + len(capBytes) + + 4 + len(reqBytes) + + 4 + len(paramsBytes) + + buf := make([]byte, size) + offset := 0 + + copy(buf[offset:], []byte(BYOCJobSigV1Prefix)) + offset += 16 + + binary.BigEndian.PutUint32(buf[offset:], uint32(job.TimeoutSeconds)) + offset += 4 + + binary.BigEndian.PutUint32(buf[offset:], uint32(len(idBytes))) + offset += 4 + copy(buf[offset:], idBytes) + offset += len(idBytes) + + binary.BigEndian.PutUint32(buf[offset:], uint32(len(capBytes))) + offset += 4 + copy(buf[offset:], capBytes) + offset += len(capBytes) + + binary.BigEndian.PutUint32(buf[offset:], uint32(len(reqBytes))) + offset += 4 + copy(buf[offset:], reqBytes) + offset += len(reqBytes) + + binary.BigEndian.PutUint32(buf[offset:], uint32(len(paramsBytes))) + offset += 4 + copy(buf[offset:], paramsBytes) + + return buf +} diff --git a/clog/clog.go b/clog/clog.go index 9ebb59a758..04b8b56989 100644 --- a/clog/clog.go +++ b/clog/clog.go @@ -197,6 +197,23 @@ func infof(ctx context.Context, lastErr bool, publicLog bool, format string, arg // Example: Info(ctx, "hello", "key1", value1, "key2", value2) // This will log: "hello key1=value1 key2=value2" func Info(ctx context.Context, msg string, keyvals ...interface{}) { + msg = formatKeyvals(msg, keyvals...) + infof(ctx, false, false, msg) +} + +// Warning logs a warning message with key-value pairs in a slog-like style. +func Warning(ctx context.Context, msg string, keyvals ...interface{}) { + msg = formatKeyvals(msg, keyvals...) + Warningf(ctx, "%s", msg) +} + +// Error logs an error message with key-value pairs in a slog-like style. +func Error(ctx context.Context, msg string, keyvals ...interface{}) { + msg = formatKeyvals(msg, keyvals...) + Errorf(ctx, "%s", msg) +} + +func formatKeyvals(msg string, keyvals ...interface{}) string { if len(keyvals)%2 != 0 { keyvals = append(keyvals[:len(keyvals)-1], "MISSING", keyvals[len(keyvals)-1]) } @@ -225,7 +242,7 @@ func Info(ctx context.Context, msg string, keyvals ...interface{}) { } } - infof(ctx, false, false, sb.String()) + return sb.String() } // V returns a Verbose instance for conditional logging at the specified level diff --git a/cmd/livepeer/starter/flags.go b/cmd/livepeer/starter/flags.go index aef859727c..1fff92adfa 100644 --- a/cmd/livepeer/starter/flags.go +++ b/cmd/livepeer/starter/flags.go @@ -56,6 +56,9 @@ func NewLivepeerConfig(fs *flag.FlagSet) LivepeerConfig { // AI: cfg.AIServiceRegistry = fs.Bool("aiServiceRegistry", *cfg.AIServiceRegistry, "Set to true to use an AI ServiceRegistry contract address") cfg.AIWorker = fs.Bool("aiWorker", *cfg.AIWorker, "Set to true to run an AI worker") + cfg.AIServerless = fs.Bool("aiServerless", *cfg.AIServerless, "Set to true to use serverless AI worker") + cfg.UseLiveRunners = fs.Bool("useLiveRunners", *cfg.UseLiveRunners, "Set to true to use LiveRunner-backed runners for supported pipelines") + cfg.LiveRunnerConfig = fs.String("liveRunnerConfig", *cfg.LiveRunnerConfig, "Path to JSON config file for static LiveRunner registration") cfg.AIModels = fs.String("aiModels", *cfg.AIModels, "Set models (pipeline:model_id) for AI worker to load upon initialization") cfg.AIModelsDir = fs.String("aiModelsDir", *cfg.AIModelsDir, "Set directory where AI model weights are stored") cfg.AIRunnerImage = fs.String("aiRunnerImage", *cfg.AIRunnerImage, "[Deprecated] Specify the base Docker image for the AI runner. Example: livepeer/ai-runner:0.0.1. Use -aiRunnerImageOverrides instead.") @@ -67,6 +70,7 @@ func NewLivepeerConfig(fs *flag.FlagSet) LivepeerConfig { // Live AI: cfg.MediaMTXApiPassword = fs.String("mediaMTXApiPassword", "", "HTTP basic auth password for MediaMTX API requests") + cfg.LiveRunnerAddr = fs.String("liveRunnerAddr", *cfg.LiveRunnerAddr, "Base URL used by live runners for heartbeat, control-plane callbacks, and internal trickle channels. Must be a full URL such as http://go-livepeer:8935") cfg.LiveAITrickleHostForRunner = fs.String("liveAITrickleHostForRunner", "", "Trickle Host used by AI Runner; It's used to overwrite the publicly available Trickle Host") cfg.LiveAIAuthApiKey = fs.String("liveAIAuthApiKey", "", "API key to use for Live AI authentication requests") cfg.LiveAIHeartbeatURL = fs.String("liveAIHeartbeatURL", "", "Base URL for Live AI heartbeat requests") @@ -145,6 +149,7 @@ func NewLivepeerConfig(fs *flag.FlagSet) LivepeerConfig { cfg.RemoteSignerWebhookHeaders = fs.String("remoteSignerWebhookHeaders", *cfg.RemoteSignerWebhookHeaders, "Map of headers to use for remote signer webhook requests. e.g. 'header:val,header2:val2'") cfg.RemoteSignerAllowNoAuth = fs.Bool("remoteSignerAllowNoAuth", *cfg.RemoteSignerAllowNoAuth, "Allow an unauthenticated remote signer on a public -httpAddr (no webhook). UNSAFE: signs payments from this node's deposit for any reachable caller; restrict access externally (proxy/private network).") cfg.RemoteDiscovery = fs.Bool("remoteDiscovery", *cfg.RemoteDiscovery, "Enable orchestrator discovery on remote signers") + cfg.ByocPerCapPricing = fs.Bool("byocPerCapPricing", *cfg.ByocPerCapPricing, "Remote signer: resolve BYOC live-payment fee from the orchestrator's per-capability CapabilitiesPrices (keyed on the request capability) instead of the flat base price. Default OFF (byte-identical to base-price behavior).") // Gateway metrics cfg.KafkaBootstrapServers = fs.String("kafkaBootstrapServers", *cfg.KafkaBootstrapServers, "URL of Kafka Bootstrap Servers") diff --git a/cmd/livepeer/starter/starter.go b/cmd/livepeer/starter/starter.go index 46e638c521..2a00ad7d35 100755 --- a/cmd/livepeer/starter/starter.go +++ b/cmd/livepeer/starter/starter.go @@ -28,6 +28,7 @@ import ( "github.com/ethereum/go-ethereum/params" "github.com/ethereum/go-ethereum/rpc" "github.com/golang/glog" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/ai/worker" "github.com/livepeer/go-livepeer/build" "github.com/livepeer/go-livepeer/common" @@ -97,6 +98,9 @@ type LivepeerConfig struct { Transcoder *bool AIServiceRegistry *bool AIWorker *bool + AIServerless *bool + UseLiveRunners *bool + LiveRunnerConfig *string Gateway *bool Broadcaster *bool OrchSecret *string @@ -164,6 +168,7 @@ type LivepeerConfig struct { AuthWebhookURL *string LiveAIAuthWebhookURL *string LiveAITrickleHostForRunner *string + LiveRunnerAddr *string OrchWebhookURL *string OrchBlacklist *string OrchMinLivepeerVersion *string @@ -175,6 +180,7 @@ type LivepeerConfig struct { RemoteSignerWebhookHeaders *string RemoteSignerAllowNoAuth *bool RemoteDiscovery *bool + ByocPerCapPricing *bool AIRunnerImage *string AIRunnerImageOverrides *string AIVerboseLogs *bool @@ -237,6 +243,9 @@ func DefaultLivepeerConfig() LivepeerConfig { // AI: defaultAIServiceRegistry := false defaultAIWorker := false + defaultAIServerless := false + defaultUseLiveRunners := false + defaultLiveRunnerConfig := "" defaultAIModels := "" defaultAIModelsDir := "" defaultAIRunnerImage := "livepeer/ai-runner:latest" @@ -246,6 +255,7 @@ func DefaultLivepeerConfig() LivepeerConfig { defaultAIMinRunnerVersion := "[]" defaultAIRunnerImageOverrides := "" defaultLiveAIAuthWebhookURL := "" + defaultLiveRunnerAddr := "" defaultLivePaymentInterval := 5 * time.Second defaultLiveOutSegmentTimeout := 0 * time.Second defaultGatewayHost := "" @@ -317,6 +327,7 @@ func DefaultLivepeerConfig() LivepeerConfig { defaultRemoteSignerWebhookHeaders := "" defaultRemoteSignerAllowNoAuth := false defaultRemoteDiscovery := false + defaultByocPerCapPricing := false // Gateway logs defaultKafkaBootstrapServers := "" @@ -363,6 +374,9 @@ func DefaultLivepeerConfig() LivepeerConfig { // AI: AIServiceRegistry: &defaultAIServiceRegistry, AIWorker: &defaultAIWorker, + AIServerless: &defaultAIServerless, + UseLiveRunners: &defaultUseLiveRunners, + LiveRunnerConfig: &defaultLiveRunnerConfig, AIModels: &defaultAIModels, AIModelsDir: &defaultAIModelsDir, AIRunnerImage: &defaultAIRunnerImage, @@ -372,6 +386,7 @@ func DefaultLivepeerConfig() LivepeerConfig { AIMinRunnerVersion: &defaultAIMinRunnerVersion, AIRunnerImageOverrides: &defaultAIRunnerImageOverrides, LiveAIAuthWebhookURL: &defaultLiveAIAuthWebhookURL, + LiveRunnerAddr: &defaultLiveRunnerAddr, LivePaymentInterval: &defaultLivePaymentInterval, LiveOutSegmentTimeout: &defaultLiveOutSegmentTimeout, GatewayHost: &defaultGatewayHost, @@ -445,6 +460,7 @@ func DefaultLivepeerConfig() LivepeerConfig { RemoteSignerWebhookHeaders: &defaultRemoteSignerWebhookHeaders, RemoteSignerAllowNoAuth: &defaultRemoteSignerAllowNoAuth, RemoteDiscovery: &defaultRemoteDiscovery, + ByocPerCapPricing: &defaultByocPerCapPricing, // Gateway logs KafkaBootstrapServers: &defaultKafkaBootstrapServers, @@ -1319,6 +1335,15 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { var aiCaps []core.Capability capabilityConstraints := make(core.PerCapabilityConstraints) + var aiModelConfigs []core.AIModelConfig + + if *cfg.AIModels != "" { + aiModelConfigs, err = core.ParseAIModelConfigs(*cfg.AIModels) + if err != nil { + glog.Errorf("Error parsing -aiModels: %v", err) + return + } + } if *cfg.AIWorker { gpus := []string{} @@ -1378,10 +1403,28 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { } } - n.AIWorker, err = worker.NewWorker(imageOverrides, *cfg.AIVerboseLogs, gpus, modelsDir, containerCreatorID) - if err != nil { - glog.Errorf("Error starting AI worker: %v", err) - return + if *cfg.AIServerless { + if len(aiModelConfigs) != 1 { + glog.Errorf("Serverless requires exactly one AI model config, got %d", len(aiModelConfigs)) + return + } + config := aiModelConfigs[0] + if config.Pipeline != "live-video-to-video" || config.ModelID != "scope" { + glog.Errorf("Serverless only supports live-video-to-video/scope, got %s/%s", config.Pipeline, config.ModelID) + return + } + + n.AIWorker, err = worker.NewServerlessWorker(strings.TrimSpace(config.URL), config.Capacity) + if err != nil { + glog.Errorf("Error starting Serverless AI worker: %v", err) + return + } + } else { + n.AIWorker, err = worker.NewWorker(imageOverrides, *cfg.AIVerboseLogs, gpus, modelsDir, containerCreatorID) + if err != nil { + glog.Errorf("Error starting AI worker: %v", err) + return + } } defer func() { @@ -1397,13 +1440,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { } if *cfg.AIModels != "" { - configs, err := core.ParseAIModelConfigs(*cfg.AIModels) - if err != nil { - glog.Errorf("Error parsing -aiModels: %v", err) - return - } - - for _, config := range configs { + for _, config := range aiModelConfigs { pipelineCap, err := core.PipelineToCapability(config.Pipeline) if err != nil { panic(fmt.Errorf("Pipeline is not valid capability: %v\n", config.Pipeline)) @@ -1411,6 +1448,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { if *cfg.AIWorker { modelConstraint := &core.ModelConstraint{Warm: config.Warm, Capacity: 1} modelsCount := 1 + if config.Capacity != 0 { if config.URL == "" { // Use multiple same configs if External Container is not used and capacity is set @@ -1846,6 +1884,9 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { if cfg.RemoteDiscovery != nil { n.RemoteDiscovery = *cfg.RemoteDiscovery } + if cfg.ByocPerCapPricing != nil { + n.ByocPerCapPricing = *cfg.ByocPerCapPricing + } if cfg.LiveAIHeartbeatHeaders != nil { n.LiveAIHeartbeatHeaders = parseHeaderMap(*cfg.LiveAIHeartbeatHeaders) } @@ -1854,6 +1895,12 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { if cfg.LiveAITrickleHostForRunner != nil { n.LiveAITrickleHostForRunner = *cfg.LiveAITrickleHostForRunner } + if cfg.LiveRunnerAddr != nil && *cfg.LiveRunnerAddr != "" { + n.LiveRunnerAddr, err = parseLiveRunnerAddr(*cfg.LiveRunnerAddr) + if err != nil { + glog.Exitf("invalid -liveRunnerAddr: %v", err) + } + } if cfg.LiveAICapRefreshModels != nil && *cfg.LiveAICapRefreshModels != "" { glog.Warningf("The -liveAICapRefreshModels flag is deprecated, capacity is now available for all models, use -liveAICapReportInterval to set the interval for reporting capacity metrics") } @@ -1935,6 +1982,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { DiscoveryTimeout: *cfg.DiscoveryTimeout, LiveAICapReportInterval: *cfg.LiveAICapReportInterval, IgnoreCapacityCheck: true, + UseDiscoveryEndpoint: true, }.New() if err != nil { exit("Could not create orchestrator pool with DB cache: %v", err) @@ -1984,6 +2032,30 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { orch := core.NewOrchestrator(s.LivepeerNode, timeWatcher) + if *cfg.UseLiveRunners || *cfg.LiveRunnerConfig != "" { + if n.OrchSecret == "" && *cfg.LiveRunnerConfig == "" { + glog.Exit("running with -useLiveRunners requires -orchSecret") + } + n.LiveRunnerManager = runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{ + Host: liveRunnerHost{RunnerHost: orch, LivepeerNode: n}, + Onchain: *cfg.Network != "offchain", + }) + if n.OrchSecret == "" { + glog.Warning("No -orchSecret configured; dynamic LiveRunner heartbeat registration is disabled") + } + if *cfg.LiveRunnerConfig != "" { + configJSON, err := os.ReadFile(*cfg.LiveRunnerConfig) + if err != nil { + glog.Exitf("error reading -liveRunnerConfig: %v", err) + } + registration, err := n.LiveRunnerManager.(*runner.LiveRunnerRegistry).RegisterStaticRunnersJSON(configJSON) + if err != nil { + glog.Exitf("error registering -liveRunnerConfig: %v", err) + } + glog.Infof("Registered %d static live runners from %s", len(registration.Runners), *cfg.LiveRunnerConfig) + } + } + go func() { err = server.StartTranscodeServer(orch, *cfg.HttpAddr, s.HTTPMux, n.WorkDir, n.TranscoderManager != nil, n.AIWorkerManager != nil, n) if err != nil { @@ -1997,10 +2069,20 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) { // check whether or not the orchestrator is available if *cfg.TestOrchAvail && doingWork { time.Sleep(2 * time.Second) - orchAvail := server.CheckOrchestratorAvailability(orch) + var ( + checkName string + orchAvail bool + ) + if *cfg.UseLiveRunners || *cfg.LiveRunnerConfig != "" { + checkName = "discovery" + orchAvail = server.CheckOrchestratorDiscoveryAvailability(orch) + } else { + checkName = "grpc ping" + orchAvail = server.CheckOrchestratorAvailability(orch) + } if !orchAvail { // shut down orchestrator - glog.Infof("Orchestrator not available at %v (%v); shutting down", orch.ServiceURI(), *cfg.HttpAddr) + glog.Infof("Orchestrator not available at %v (%v) via %s check; shutting down", orch.ServiceURI(), *cfg.HttpAddr, checkName) tc <- struct{}{} } } @@ -2143,7 +2225,11 @@ func parseHeaderMap(raw string) map[string]string { for _, header := range strings.Split(raw, ",") { parts := strings.SplitN(header, ":", 2) if len(parts) == 2 { - headers[parts[0]] = parts[1] + key := strings.TrimSpace(parts[0]) + value := strings.TrimSpace(parts[1]) + if key != "" { + headers[key] = value + } } } return headers @@ -2177,7 +2263,10 @@ func getServiceURI(n *core.LivepeerNode, serviceAddr string) (*url.URL, error) { // special value to signal this node is not to be used for work return url.Parse("") } - return url.ParseRequestURI("https://" + serviceAddr) + if !strings.HasPrefix(serviceAddr, "http://") && !strings.HasPrefix(serviceAddr, "https://") { + serviceAddr = "https://" + serviceAddr + } + return url.ParseRequestURI(serviceAddr) } // Infer address @@ -2504,3 +2593,30 @@ func exit(msg string, args ...any) { glog.Errorf(msg, args...) os.Exit(2) } + +type liveRunnerHost struct { + runner.RunnerHost + *core.LivepeerNode +} + +func (h liveRunnerHost) LiveRunnerURI() *url.URL { + if h.LivepeerNode != nil && h.LivepeerNode.LiveRunnerAddr != nil { + v := *h.LivepeerNode.LiveRunnerAddr + return &v + } + return h.RunnerHost.ServiceURI() +} + +func parseLiveRunnerAddr(addr string) (*url.URL, error) { + parsed, err := url.ParseRequestURI(addr) + if err != nil { + return nil, err + } + if !parsed.IsAbs() || parsed.Host == "" { + return nil, fmt.Errorf("must be an absolute URL") + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return nil, fmt.Errorf("scheme must be http or https") + } + return parsed, nil +} diff --git a/cmd/livepeer/starter/starter_test.go b/cmd/livepeer/starter/starter_test.go index 705426eb32..c1e9dec083 100644 --- a/cmd/livepeer/starter/starter_test.go +++ b/cmd/livepeer/starter/starter_test.go @@ -10,6 +10,7 @@ import ( "testing" ethcommon "github.com/ethereum/go-ethereum/common" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" "github.com/livepeer/go-livepeer/eth" @@ -89,6 +90,44 @@ func TestIsLocalURL(t *testing.T) { assert.False(isLocal) } +func TestGetServiceURIServiceAddrScheme(t *testing.T) { + uri, err := getServiceURI(nil, "127.0.0.1:8935") + require.NoError(t, err) + require.Equal(t, "https://127.0.0.1:8935", uri.String()) + + uri, err = getServiceURI(nil, "http://127.0.0.1:8935") + require.NoError(t, err) + require.Equal(t, "http://127.0.0.1:8935", uri.String()) + + uri, err = getServiceURI(nil, "https://orch.example.com:443") + require.NoError(t, err) + require.Equal(t, "https://orch.example.com:443", uri.String()) + + uri, err = getServiceURI(nil, "gopher://orch.example.com:443") + require.NoError(t, err) + require.Equal(t, "https://gopher://orch.example.com:443", uri.String()) + + uri, err = getServiceURI(nil, "none") + require.NoError(t, err) + require.Equal(t, "", uri.String()) +} + +func TestParseLiveRunnerAddr(t *testing.T) { + uri, err := parseLiveRunnerAddr("http://go-livepeer:8935") + require.NoError(t, err) + require.Equal(t, "http://go-livepeer:8935", uri.String()) + + uri, err = parseLiveRunnerAddr("https://public.example.com") + require.NoError(t, err) + require.Equal(t, "https://public.example.com", uri.String()) + + _, err = parseLiveRunnerAddr("go-livepeer:8935") + require.Error(t, err) + + _, err = parseLiveRunnerAddr("ftp://go-livepeer:8935") + require.Error(t, err) +} + func TestParseGetGatewayPrices(t *testing.T) { assert := assert.New(t) @@ -377,11 +416,14 @@ func TestPrintConfigRedaction(t *testing.T) { func TestParseHeaderMap(t *testing.T) { require := require.New(t) - headers := parseHeaderMap("Authorization:Bearer abc,X-API-Key:secret,invalid") + headers := parseHeaderMap("Authorization: Bearer abc, X-Webhook-Secret: webhook-secret, X-API-Key: secret,invalid, :missing-key") require.Equal("Bearer abc", headers["Authorization"]) + require.Equal("webhook-secret", headers["X-Webhook-Secret"]) require.Equal("secret", headers["X-API-Key"]) _, exists := headers["invalid"] require.False(exists) + _, exists = headers[""] + require.False(exists) } func TestNewLivepeerConfig_RemoteSignerWebhookFlags(t *testing.T) { @@ -400,6 +442,56 @@ func TestNewLivepeerConfig_RemoteSignerWebhookFlags(t *testing.T) { require.Equal("Authorization:Bearer abc,X-API-Key:secret", *cfg.RemoteSignerWebhookHeaders) } +func TestNewLivepeerConfig_UseLiveRunnersFlag(t *testing.T) { + require := require.New(t) + + fs := flag.NewFlagSet("livepeer-test", flag.ContinueOnError) + cfg := NewLivepeerConfig(fs) + require.NoError(fs.Parse([]string{"-useLiveRunners"})) + require.True(*cfg.UseLiveRunners) +} + +func TestNewLivepeerConfig_LiveRunnerConfigFlag(t *testing.T) { + require := require.New(t) + + fs := flag.NewFlagSet("livepeer-test", flag.ContinueOnError) + cfg := NewLivepeerConfig(fs) + require.NoError(fs.Parse([]string{"-liveRunnerConfig", "/tmp/runners.json"})) + require.Equal("/tmp/runners.json", *cfg.LiveRunnerConfig) +} + +func TestNewLivepeerConfig_LiveRunnerAddrFlag(t *testing.T) { + require := require.New(t) + + fs := flag.NewFlagSet("livepeer-test", flag.ContinueOnError) + cfg := NewLivepeerConfig(fs) + require.NoError(fs.Parse([]string{"-liveRunnerAddr", "http://go-livepeer:8935"})) + require.Equal("http://go-livepeer:8935", *cfg.LiveRunnerAddr) +} + +func TestNewLivepeerConfig_UseLiveWorkersFlagRemoved(t *testing.T) { + fs := flag.NewFlagSet("livepeer-test", flag.ContinueOnError) + NewLivepeerConfig(fs) + require.Error(t, fs.Parse([]string{"-useLiveWorkers"})) +} + +func TestLiveRunnerOrchSecretLiteralBehavior(t *testing.T) { + secret, err := common.ReadFromFile("literal-secret") + require.Error(t, err) + require.Equal(t, "literal-secret", secret) + require.NotEmpty(t, secret) +} + +func TestLiveRunnerManagerConstruction(t *testing.T) { + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{}) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + require.NotNil(t, node.LiveRunnerManager) +} + // Helper struct to capture output for testing type testWriter struct { buf *[]byte diff --git a/common/types.go b/common/types.go index 1fef7d7886..37368c33c5 100644 --- a/common/types.go +++ b/common/types.go @@ -177,4 +177,5 @@ type OrchNetworkCapabilities struct { PriceInfo *net.PriceInfo `json:"price_info"` CapabilitiesPrices []*net.PriceInfo `json:"capabilities_prices"` Hardware []*net.HardwareInformation `json:"hardware"` + Discovery json.RawMessage `json:"-"` } diff --git a/core/accounting.go b/core/accounting.go index 213389e933..b166833632 100644 --- a/core/accounting.go +++ b/core/accounting.go @@ -104,6 +104,28 @@ func (a *AddressBalances) Balance(addr ethcommon.Address, id ManifestID) *big.Ra return a.balancesForAddr(addr).Balance(id) } +// FixedPrice retrieves the fixed price for an address' ManifestID. +func (a *AddressBalances) FixedPrice(addr ethcommon.Address, id ManifestID) *big.Rat { + a.mtx.Lock() + balances := a.balances[addr] + a.mtx.Unlock() + if balances == nil { + return nil + } + return balances.FixedPrice(id) +} + +// SetFixedPrice sets the fixed price for an address' ManifestID. +func (a *AddressBalances) SetFixedPrice(addr ethcommon.Address, id ManifestID, fixedPrice *big.Rat) { + a.mtx.Lock() + balances := a.balances[addr] + a.mtx.Unlock() + if balances == nil { + return + } + balances.SetFixedPrice(id, fixedPrice) +} + // StopCleanup stops the cleanup loop for all balances func (a *AddressBalances) StopCleanup() { a.mtx.Lock() diff --git a/core/capabilities.go b/core/capabilities.go index e1ecd62576..838ede18f5 100644 --- a/core/capabilities.go +++ b/core/capabilities.go @@ -5,6 +5,8 @@ import ( "encoding/json" "errors" "fmt" + "sort" + "strings" "sync" "github.com/Masterminds/semver/v3" @@ -791,6 +793,80 @@ func (bcast *Capabilities) PerCapability() PerCapabilityConstraints { return nil } +// ModelIDForCapability returns the constrained model ID for the given capability. +// When multiple models are constrained it returns the lexicographically first one +// so the result is deterministic. Returns "" when no model is constrained. +func (bcast *Capabilities) ModelIDForCapability(cap Capability) string { + if bcast == nil { + return "" + } + capConstraints, ok := bcast.constraints.perCapability[cap] + if !ok || capConstraints == nil { + return "" + } + modelIDs := make([]string, 0, len(capConstraints.Models)) + for modelID := range capConstraints.Models { + if modelID != "" { + modelIDs = append(modelIDs, modelID) + } + } + if len(modelIDs) == 0 { + return "" + } + sort.Strings(modelIDs) + return modelIDs[0] +} + +// CapabilityToPipeline converts a capability to a pipeline slug. +// Returns "" when the capability is unknown. +func CapabilityToPipeline(cap Capability) string { + capName, ok := CapabilityNameLookup[cap] + if !ok { + return "" + } + return strings.ReplaceAll(strings.ToLower(capName), " ", "-") +} + +// ConstrainedPipelineModelID returns a deterministic pipeline/model_id pair +// from per-capability constraints. When multiple capabilities have constrained +// models, it returns the lexicographically first pipeline. +func (bcast *Capabilities) ConstrainedPipelineModelID() (string, string) { + if bcast == nil { + return "", "" + } + + type capPipeline struct { + capability Capability + pipeline string + } + + candidates := []capPipeline{} + for capability := range bcast.constraints.perCapability { + modelID := bcast.ModelIDForCapability(capability) + if modelID == "" { + continue + } + pipeline := CapabilityToPipeline(capability) + if pipeline == "" { + continue + } + candidates = append(candidates, capPipeline{ + capability: capability, + pipeline: pipeline, + }) + } + if len(candidates) == 0 { + return "", "" + } + + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].pipeline < candidates[j].pipeline + }) + + selected := candidates[0] + return selected.pipeline, bcast.ModelIDForCapability(selected.capability) +} + func (bcast *Capabilities) SetMinVersionConstraint(minVersionConstraint string) { if bcast != nil { bcast.constraints.minVersion = minVersionConstraint diff --git a/core/capabilities_test.go b/core/capabilities_test.go index 59f45b0582..db3f0a6e61 100644 --- a/core/capabilities_test.go +++ b/core/capabilities_test.go @@ -751,6 +751,110 @@ func TestMinRunnerVersion(t *testing.T) { assert.Equal("", c.MinRunnerVersionConstraint(Capability_LiveVideoToVideo, "other")) } +func TestModelIDForCapability(t *testing.T) { + assert := assert.New(t) + + // nil receiver is safe + var nilCaps *Capabilities + assert.Equal("", nilCaps.ModelIDForCapability(Capability_LiveVideoToVideo)) + + // no constraints for the capability + c := NewCapabilities([]Capability{Capability_LiveVideoToVideo}, nil) + assert.Equal("", c.ModelIDForCapability(Capability_LiveVideoToVideo)) + + // single constrained model + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_LiveVideoToVideo: &CapabilityConstraints{ + Models: ModelConstraints{"streamdiffusion": &ModelConstraint{Warm: true}}, + }, + }) + assert.Equal("streamdiffusion", c.ModelIDForCapability(Capability_LiveVideoToVideo)) + + // different capability has no constraint + assert.Equal("", c.ModelIDForCapability(Capability_TextToImage)) + + // empty model IDs are ignored + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_LiveVideoToVideo: &CapabilityConstraints{ + Models: ModelConstraints{"": &ModelConstraint{}}, + }, + }) + assert.Equal("", c.ModelIDForCapability(Capability_LiveVideoToVideo)) + + // one real model plus empty key resolves to the real model + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_LiveVideoToVideo: &CapabilityConstraints{ + Models: ModelConstraints{ + "": &ModelConstraint{}, + "streamdiffusion": &ModelConstraint{Warm: true}, + }, + }, + }) + assert.Equal("streamdiffusion", c.ModelIDForCapability(Capability_LiveVideoToVideo)) + + // multiple models resolve to the lexicographically first ID + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_LiveVideoToVideo: &CapabilityConstraints{ + Models: ModelConstraints{ + "comfyui": &ModelConstraint{}, + "streamdiffusion": &ModelConstraint{}, + "noop": &ModelConstraint{}, + }, + }, + }) + assert.Equal("comfyui", c.ModelIDForCapability(Capability_LiveVideoToVideo)) +} + +func TestCapabilityToPipeline(t *testing.T) { + assert := assert.New(t) + + assert.Equal("live-video-to-video", CapabilityToPipeline(Capability_LiveVideoToVideo)) + assert.Equal("text-to-image", CapabilityToPipeline(Capability_TextToImage)) + assert.Equal("", CapabilityToPipeline(Capability(-999))) +} + +func TestConstrainedPipelineModelID(t *testing.T) { + assert := assert.New(t) + + var nilCaps *Capabilities + pipeline, modelID := nilCaps.ConstrainedPipelineModelID() + assert.Equal("", pipeline) + assert.Equal("", modelID) + + c := NewCapabilities([]Capability{Capability_LiveVideoToVideo, Capability_TextToImage}, nil) + pipeline, modelID = c.ConstrainedPipelineModelID() + assert.Equal("", pipeline) + assert.Equal("", modelID) + + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_TextToImage: &CapabilityConstraints{ + Models: ModelConstraints{ + "stabilityai/sd-turbo": &ModelConstraint{Warm: true}, + }, + }, + }) + pipeline, modelID = c.ConstrainedPipelineModelID() + assert.Equal("text-to-image", pipeline) + assert.Equal("stabilityai/sd-turbo", modelID) + + // Multiple constrained capabilities resolve deterministically. + c.SetPerCapabilityConstraints(PerCapabilityConstraints{ + Capability_LiveVideoToVideo: &CapabilityConstraints{ + Models: ModelConstraints{ + "streamdiffusion": &ModelConstraint{}, + }, + }, + Capability_AudioToText: &CapabilityConstraints{ + Models: ModelConstraints{ + "whisper-large": &ModelConstraint{}, + }, + }, + }) + pipeline, modelID = c.ConstrainedPipelineModelID() + assert.Equal("audio-to-text", pipeline) + assert.Equal("whisper-large", modelID) +} + func (c *Constraints) addCapabilityConstraints(cap Capability, constraint CapabilityConstraints) { // the capability should be added by AddCapacity for modelID, modelConstraint := range constraint.Models { diff --git a/core/livepeernode.go b/core/livepeernode.go index b6971b6f96..0ace42dd55 100644 --- a/core/livepeernode.go +++ b/core/livepeernode.go @@ -128,6 +128,8 @@ type LivepeerNode struct { AIWorkerManager *RemoteAIWorkerManager AIProcesssingRetryTimeout time.Duration + LiveRunnerManager any // NB: kludge to avoid ai/runner circular dependency + // Transcoder public fields SegmentChans map[ManifestID]SegmentChan Recipient pm.Recipient @@ -156,6 +158,7 @@ type LivepeerNode struct { RemoteEthAddr ethcommon.Address // eth address of the remote signer InfoSig []byte // sig over eth address for the OrchestratorInfo request RemoteDiscovery bool // expose remote discovery endpoint when enabled + ByocPerCapPricing bool // resolve BYOC fee from per-capability CapabilitiesPrices instead of base price (remote signer; default OFF) // Thread safety for config fields mu sync.RWMutex @@ -174,6 +177,7 @@ type LivepeerNode struct { LivePipelines map[string]*LivePipeline LiveMu *sync.RWMutex + LiveRunnerAddr *url.URL MediaMTXApiPassword string LiveAITrickleHostForRunner string LiveAIAuthWebhookURL *url.URL diff --git a/core/orchestrator.go b/core/orchestrator.go index fd33eddc5a..0d4d8417a2 100644 --- a/core/orchestrator.go +++ b/core/orchestrator.go @@ -56,6 +56,10 @@ func (orch *orchestrator) ServiceURI() *url.URL { return orch.node.GetServiceURI() } +func (orch *orchestrator) LiveRunnerURI() *url.URL { + return orch.ServiceURI() +} + func (orch *orchestrator) Sign(msg []byte) ([]byte, error) { if orch.node == nil || orch.node.Eth == nil { return []byte{}, nil @@ -78,6 +82,10 @@ func (orch *orchestrator) TranscoderSecret() string { return orch.node.OrchSecret } +func (orch *orchestrator) RegistrationSecret() string { + return orch.TranscoderSecret() +} + func (orch *orchestrator) CheckCapacity(mid ManifestID) error { orch.node.segmentMutex.RLock() defer orch.node.segmentMutex.RUnlock() @@ -383,11 +391,9 @@ func (orch *orchestrator) PriceInfoForCaps(sender ethcommon.Address, manifestID func (orch *orchestrator) priceInfo(sender ethcommon.Address, manifestID ManifestID, caps *net.Capabilities) (*big.Rat, error) { // If there is already a fixed price for the given session, use this price if manifestID != "" { - if balances, ok := orch.node.Balances.balances[sender]; ok { - fixedPrice := balances.FixedPrice(manifestID) - if fixedPrice != nil { - return fixedPrice, nil - } + fixedPrice := orch.node.Balances.FixedPrice(sender, manifestID) + if fixedPrice != nil { + return fixedPrice, nil } } @@ -532,11 +538,9 @@ func (orch *orchestrator) setFixedPricePerSession(sender ethcommon.Address, mani glog.Warning("Node balances are not initialized") return } - if balances, ok := orch.node.Balances.balances[sender]; ok { - if balances.FixedPrice(manifestID) == nil { - balances.SetFixedPrice(manifestID, priceInfoRat) - glog.V(6).Infof("Setting fixed price=%v for session=%v", priceInfoRat, manifestID) - } + if orch.node.Balances.FixedPrice(sender, manifestID) == nil { + orch.node.Balances.SetFixedPrice(sender, manifestID, priceInfoRat) + glog.V(6).Infof("Setting fixed price=%v for session=%v", priceInfoRat, manifestID) } } diff --git a/discovery/db_discovery.go b/discovery/db_discovery.go index 8e85cc8c5b..4f4e5fbc76 100644 --- a/discovery/db_discovery.go +++ b/discovery/db_discovery.go @@ -2,8 +2,11 @@ package discovery import ( "context" + "encoding/json" "fmt" + "io" "math/big" + "net/http" "net/url" "strings" "time" @@ -22,6 +25,10 @@ import ( "github.com/golang/glog" ) +const orchestratorEndpointDiscoveryMaxBytes = 1 << 20 + +var orchestratorEndpointDiscoveryTimeout = 2 * time.Second + type ticketParamsValidator interface { ValidateTicketParams(ticketParams *pm.TicketParams) error } @@ -35,9 +42,17 @@ type DBOrchestratorPoolCache struct { orchBlacklist []string discoveryTimeout time.Duration ignoreCapacityCheck bool + useDiscoveryEndpoint bool node *core.LivepeerNode } +type orchPollingInfo struct { + level int + orchInfo *net.OrchestratorInfo + dbOrch *common.DBOrch + discovery json.RawMessage +} + func NewDBOrchestratorPoolCache(ctx context.Context, node *core.LivepeerNode, rm common.RoundsManager, orchBlacklist []string, discoveryTimeout time.Duration, liveAICapReportInterval time.Duration) (*DBOrchestratorPoolCache, error) { return DBOrchestratorPoolCacheConfig{ Ctx: ctx, @@ -57,6 +72,7 @@ type DBOrchestratorPoolCacheConfig struct { DiscoveryTimeout time.Duration LiveAICapReportInterval time.Duration IgnoreCapacityCheck bool + UseDiscoveryEndpoint bool } func (cfg DBOrchestratorPoolCacheConfig) New() (*DBOrchestratorPoolCache, error) { @@ -74,6 +90,7 @@ func (cfg DBOrchestratorPoolCacheConfig) New() (*DBOrchestratorPoolCache, error) orchBlacklist: cfg.OrchBlacklist, discoveryTimeout: cfg.DiscoveryTimeout, ignoreCapacityCheck: cfg.IgnoreCapacityCheck, + useDiscoveryEndpoint: cfg.UseDiscoveryEndpoint, node: node, } @@ -338,12 +355,6 @@ func (dbo *DBOrchestratorPoolCache) cacheOrchInfos() error { glog.Infof("Using DB orchestrator pool with %d orchestrators", len(orchs)) } - type orchPollingInfo struct { - level int - orchInfo *net.OrchestratorInfo - dbOrch *common.DBOrch - } - nodesPerOrch := dbo.bcast.ExtraNodes() // Each base orchestrator can contribute itself plus up to nodesPerOrch first-level advertised nodes. maxOrchs := len(orchs) * (nodesPerOrch + 1) @@ -367,6 +378,19 @@ func (dbo *DBOrchestratorPoolCache) cacheOrchInfos() error { errc <- fmt.Errorf("skipping orch=%v, URI not set", orch.URL.String()) return } + + var discoveryCh chan json.RawMessage + if dbo.useDiscoveryEndpoint { + discoveryCh = make(chan json.RawMessage, 1) + go func() { + discovery, err := callOrchestratorDiscovery(ctx, uri) + if err != nil { + clog.V(common.DEBUG).Infof(ctx, "unable to fetch orchestrator endpoint discovery orch=%v err=%q", uri, err) + } + discoveryCh <- discovery + }() + } + info, err := getOrchInfoRPC(ctx, dbo.bcast, uri, server.GetOrchestratorInfoParams{ IgnoreCapacityCheck: dbo.ignoreCapacityCheck, }) @@ -407,10 +431,20 @@ func (dbo *DBOrchestratorPoolCache) cacheOrchInfos() error { } } + var discovery json.RawMessage + if discoveryCh != nil { + select { + case discovery = <-discoveryCh: + case <-ctx.Done(): + clog.V(common.DEBUG).Infof(ctx, "skipping orchestrator endpoint discovery orch=%v err=%q", uri, ctx.Err()) + } + } + resc <- orchPollingInfo{ - level: level, - orchInfo: info, - dbOrch: dbOrch, + level: level, + orchInfo: info, + dbOrch: dbOrch, + discovery: discovery, } } @@ -438,7 +472,7 @@ func (dbo *DBOrchestratorPoolCache) cacheOrchInfos() error { select { case res := <-resc: //add response to network capabilities - orchNetworkCapabilities = append(orchNetworkCapabilities, orchInfoToOrchNetworkCapabilities(res.orchInfo)) + orchNetworkCapabilities = append(orchNetworkCapabilities, orchInfoToOrchNetworkCapabilities(res)) // discover newly advertised nodes. only recurse the first level. if res.level == 0 && len(res.orchInfo.GetNodes()) > 0 { @@ -578,10 +612,11 @@ func pmTicketParams(params *net.TicketParams) *pm.TicketParams { } } -func orchInfoToOrchNetworkCapabilities(info *net.OrchestratorInfo) *common.OrchNetworkCapabilities { +func orchInfoToOrchNetworkCapabilities(res orchPollingInfo) *common.OrchNetworkCapabilities { var orch common.OrchNetworkCapabilities // add orch operating information if available + info := res.orchInfo if info != nil { orch.LocalAddress = ethcommon.BytesToAddress(info.GetAddress()).Hex() orch.OrchURI = info.GetTranscoder() @@ -593,6 +628,50 @@ func orchInfoToOrchNetworkCapabilities(info *net.OrchestratorInfo) *common.OrchN orch.Address = string(ethcommon.BytesToAddress(info.TicketParams.Recipient).Hex()) } } + orch.Discovery = res.discovery return &orch } + +func callOrchestratorDiscovery(ctx context.Context, orchURI *url.URL) (json.RawMessage, error) { + if orchURI == nil { + return nil, fmt.Errorf("missing orchestrator URI") + } + if orchURI.Host == "" { + return nil, fmt.Errorf("missing host in orchestrator URI %q", orchURI.String()) + } + + discoveryURL := orchURI.JoinPath("discovery") + reqCtx, cancel := context.WithTimeout(ctx, orchestratorEndpointDiscoveryTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, discoveryURL.String(), nil) + if err != nil { + return nil, err + } + + client := &http.Client{Timeout: orchestratorEndpointDiscoveryTimeout} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return nil, fmt.Errorf("endpoint discovery returned status %d", resp.StatusCode) + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, orchestratorEndpointDiscoveryMaxBytes+1)) + if err != nil { + return nil, err + } + if len(body) > orchestratorEndpointDiscoveryMaxBytes { + return nil, fmt.Errorf("endpoint discovery response exceeds %d bytes", orchestratorEndpointDiscoveryMaxBytes) + } + + if !json.Valid(body) { + return nil, fmt.Errorf("invalid endpoint discovery JSON") + } + + return json.RawMessage(body), nil +} diff --git a/discovery/discovery_test.go b/discovery/discovery_test.go index c6706c1280..68f16fccae 100644 --- a/discovery/discovery_test.go +++ b/discovery/discovery_test.go @@ -13,6 +13,7 @@ import ( "net/url" "runtime" "strconv" + "strings" "sync" "testing" "testing/synctest" @@ -2166,6 +2167,105 @@ func TestOrchestratorPool_LatencySorting(t *testing.T) { synctest.Test(t, sync_TestOrchestratorPool_LatencySorting) } +func TestFetchOrchestratorEndpointDiscovery(t *testing.T) { + t.Run("valid entries", func(t *testing.T) { + var serverURL string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/discovery", r.URL.Path) + _, _ = fmt.Fprintf(w, `[ + {"address":"https://other.example.com","runners":[{"app":"live-video-to-video/other"}]}, + {"address":%q,"runners":[{"app":"live-video-to-video/model-a"}]} + ]`, serverURL) + })) + defer ts.Close() + serverURL = ts.URL + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.NoError(t, err) + var entries []map[string]json.RawMessage + require.NoError(t, json.Unmarshal(discovery, &entries)) + require.Len(t, entries, 2) + require.JSONEq(t, `"https://other.example.com"`, string(entries[0]["address"])) + require.JSONEq(t, fmt.Sprintf("%q", ts.URL), string(entries[1]["address"])) + require.Contains(t, string(entries[1]["runners"]), "live-video-to-video/model-a") + }) + + t.Run("non-200 is non-fatal error", func(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadGateway) + })) + defer ts.Close() + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.Error(t, err) + require.Nil(t, discovery) + }) + + t.Run("invalid JSON", func(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`not-json`)) + })) + defer ts.Close() + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.Error(t, err) + require.Nil(t, discovery) + }) + + t.Run("oversized response", func(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(strings.Repeat("x", orchestratorEndpointDiscoveryMaxBytes+1))) + })) + defer ts.Close() + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.Error(t, err) + require.Nil(t, discovery) + }) + + t.Run("entries without matching address are returned", func(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`[{"address":"https://other.example.com"}]`)) + })) + defer ts.Close() + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.NoError(t, err) + var entries []map[string]json.RawMessage + require.NoError(t, json.Unmarshal(discovery, &entries)) + require.Len(t, entries, 1) + require.JSONEq(t, `"https://other.example.com"`, string(entries[0]["address"])) + }) + + t.Run("timeout", func(t *testing.T) { + oldTimeout := orchestratorEndpointDiscoveryTimeout + orchestratorEndpointDiscoveryTimeout = time.Millisecond + defer func() { orchestratorEndpointDiscoveryTimeout = oldTimeout }() + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(20 * time.Millisecond) + _, _ = w.Write([]byte(`[]`)) + })) + defer ts.Close() + + orchURI, err := url.ParseRequestURI(ts.URL) + require.NoError(t, err) + discovery, err := callOrchestratorDiscovery(context.Background(), orchURI) + require.Error(t, err) + require.Nil(t, discovery) + }) +} + func wgWait(wg *sync.WaitGroup) bool { c := make(chan struct{}) go func() { defer close(c); wg.Wait() }() diff --git a/doc/remote-signer.md b/doc/remote-signer.md index bfb1d94cdb..dfeb2389ad 100644 --- a/doc/remote-signer.md +++ b/doc/remote-signer.md @@ -77,7 +77,7 @@ When enabled, the signer exposes: - `GET /discover-orchestrators` -The endpoint returns a list of orchestrators (`address`, `score`, `capabilities`) in a format that is compatible with the gateway's orchestrator discovery webhook. +The endpoint returns a list of orchestrators (`address`, `score`, `capabilities`, and optional `runners`) in a format that is compatible with the gateway's orchestrator discovery webhook. The `address` field is the orchestrator service address used by gateways, not the signer's Ethereum/account address. Clients can filter for orchestrators matching a given capability via the `caps` query param. This param can be repeated to retrieve orchestrators supporting any one of the given capabilities. For example: @@ -94,6 +94,10 @@ curl "http://127.0.0.1:7936/discover-orchestrators?caps=live-video-to-video/stre The remote signer periodically retrieves latest orchestrator capabilities and pricing from the network. The periodicity can be configured via the `-liveAICapReportInterval` flag with a default of 25 minutes. Orchestrators are pre-filtered for pricing: orchestrators that have a price higher than what the remote signer is configured for will not be made available via discovery. +In on-chain remote discovery mode, the signer also reads each orchestrator's `/discovery` endpoint. This lets the signer expose orchestrator service addresses learned from runner discovery. + +Runner discovery is merged by runner `url`. The first runner value wins; identical duplicates are ignored, and conflicting duplicates are logged. Exposed runners must include `price_info`; eligible runner `app` values are used for `/discover-orchestrators` capability filtering. + Currently, remote discovery can only be enabled for nodes in remote signing mode. ### Gateway node diff --git a/docker/Dockerfile b/docker/Dockerfile index efa33a4d47..1ae4ed8032 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -61,6 +61,9 @@ COPY ./install_ffmpeg.sh ./install_ffmpeg.sh ARG BUILD_TAGS ENV BUILD_TAGS=${BUILD_TAGS} +ARG GIT_REVISION=unknown +RUN echo "Building livepeer at ${GIT_REVISION}" + COPY go.mod go.sum ./ RUN go mod download diff --git a/server/ai_http.go b/server/ai_http.go index 0d61c26cd7..e75096261e 100644 --- a/server/ai_http.go +++ b/server/ai_http.go @@ -15,6 +15,7 @@ import ( "mime" "mime/multipart" "net/http" + "net/http/httputil" url2 "net/url" "strconv" "strings" @@ -24,11 +25,14 @@ import ( ethcommon "github.com/ethereum/go-ethereum/common" "github.com/getkin/kin-openapi/openapi3filter" "github.com/golang/glog" + "github.com/golang/protobuf/proto" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/ai/worker" "github.com/livepeer/go-livepeer/clog" "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" "github.com/livepeer/go-livepeer/monitor" + lpnet "github.com/livepeer/go-livepeer/net" "github.com/livepeer/go-livepeer/trickle" middleware "github.com/oapi-codegen/nethttp-middleware" ) @@ -37,6 +41,10 @@ var MaxAIRequestSize = 3000000000 // 3GB var TrickleHTTPPath = "/ai/trickle/" +const maxLiveRunnerTrickleChannelsPerRequest = 25 +const liveRunnerSenderHeader = "Livepeer-Payer-Address" +const maxScopeRequestBodySize = 1 << 20 // 1 MB + func startAIServer(lp *lphttp) error { swagger, err := worker.GetSwagger() if err != nil { @@ -61,6 +69,14 @@ func startAIServer(lp *lphttp) error { Mux: lp.transRPC, BasePath: TrickleHTTPPath, }) + if sw, ok := lp.node.AIWorker.(*worker.ServerlessWorker); ok { + sw.SetTrickleServer(lp.trickleSrv) + } + if manager, ok := lp.liveRunnerManager(); ok { + publicTrickleBaseURL := lp.orchestrator.ServiceURI().JoinPath(TrickleHTTPPath).String() + internalTrickleBaseURL := liveRunnerURI(lp.node, lp.orchestrator).JoinPath(TrickleHTTPPath).String() + manager.SetTrickleServer(lp.trickleSrv, publicTrickleBaseURL, internalTrickleBaseURL) + } lp.transRPC.Handle("/text-to-image", oapiReqValidator(aiHttpHandle(lp, jsonDecoder[worker.GenTextToImageJSONRequestBody]))) lp.transRPC.Handle("/image-to-image", oapiReqValidator(aiHttpHandle(lp, multipartDecoder[worker.GenImageToImageMultipartRequestBody]))) @@ -71,12 +87,709 @@ func startAIServer(lp *lphttp) error { lp.transRPC.Handle("/segment-anything-2", oapiReqValidator(aiHttpHandle(lp, multipartDecoder[worker.GenSegmentAnything2MultipartRequestBody]))) lp.transRPC.Handle("/image-to-text", oapiReqValidator(aiHttpHandle(lp, multipartDecoder[worker.GenImageToTextMultipartRequestBody]))) lp.transRPC.Handle("/text-to-speech", oapiReqValidator(aiHttpHandle(lp, jsonDecoder[worker.GenTextToSpeechJSONRequestBody]))) + lp.transRPC.Handle("/scope", lp.StartScope()) lp.transRPC.Handle("/live-video-to-video", oapiReqValidator(lp.StartLiveVideoToVideo())) + // Internal runner endpoints + lp.transRPC.HandleFunc("POST /runners/heartbeat", lp.LiveRunnerHeartbeat) + lp.transRPC.HandleFunc("POST /runners/{runner_id}/unregister", lp.UnregisterLiveRunner) + lp.transRPC.HandleFunc("POST /runner/{runner_id}/session/{session_id}/channels", lp.CreateLiveRunnerTrickleChannel) + lp.transRPC.HandleFunc("DELETE /runner/{runner_id}/session/{session_id}/channels", lp.DeleteLiveRunnerTrickleChannels) + lp.transRPC.HandleFunc("POST /runner/{runner_id}/session/{session_id}/stop", lp.StopLiveRunnerSessionInternal) + // Public client endpoints + lp.transRPC.HandleFunc("GET /discovery", lp.DiscoverLiveRunners) + lp.transRPC.HandleFunc("POST /apps/{runner_id}/session", lp.ReserveLiveRunnerSession) + lp.transRPC.HandleFunc("POST /apps/{runner_id}/session/{session_id}/stop", lp.StopLiveRunnerSession) + lp.transRPC.HandleFunc("POST /apps/{runner_id}/session/{session_id}/payment", lp.PaymentForLiveRunnerSession) + lp.transRPC.HandleFunc("/apps/{runner_id}/session/{session_id}/app/{app_path...}", lp.ProxyLiveRunnerSession) + lp.transRPC.HandleFunc("/apps/{runner_id}/app/{app_path...}", lp.ProxyLiveRunnerSingleShot) + // Additionally, there is the '/aiResults' endpoint registered in server/rpc.go return nil } +type liveRunnerManager interface { + Heartbeat(req runner.LiveRunnerHeartbeatRequest, auth string) (*runner.LiveRunnerHeartbeatResponse, error) + Unregister(runnerID, auth string) error + Runners() []runner.LiveRunnerDiscoveryRunner + PaymentInfo(runnerID string) (*runner.LiveRunnerPriceInfo, error) + ReserveSession(runnerID string, sessionID ...string) (string, string, error) + ReleaseSession(runnerID, sessionID string) error + RunnerMode(runnerID string) (string, error) + RunnerEndpointForSession(runnerID, sessionID string) (string, error) + SessionTokenForSession(runnerID, sessionID string) (string, error) + ValidSessionToken(runnerID, sessionID, token string) error + SetTrickleServer(srv *trickle.Server, publicBaseURL, internalBaseURL string) + CreateTrickleChannel(runnerID, sessionID, name, mimeType string) (runner.LiveRunnerTrickleChannel, error) + DeleteTrickleChannel(runnerID, sessionID, name string) error +} + +func (h *lphttp) liveRunnerManager() (liveRunnerManager, bool) { + if h.node == nil { + return nil, false + } + manager, ok := h.node.LiveRunnerManager.(liveRunnerManager) + return manager, ok +} + +func (h *lphttp) LiveRunnerHeartbeat(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + var req runner.LiveRunnerHeartbeatRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + respondWithError(w, err.Error(), http.StatusBadRequest) + return + } + resp, err := manager.Heartbeat(req, r.Header.Get("Authorization")) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + data, err := json.Marshal(resp) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + respondJsonOk(w, data) +} + +func (h *lphttp) UnregisterLiveRunner(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + if err := manager.Unregister(runnerID, r.Header.Get("Authorization")); err != nil { + respondWithLiveRunnerError(w, err) + return + } + w.WriteHeader(http.StatusNoContent) +} + +type liveRunnerSessionResponse struct { + SessionID string `json:"session_id"` + AppURL string `json:"app_url"` + ControlURL string `json:"control_url"` +} + +type liveRunnerPaymentChallengeResponse struct { + PaymentParams string `json:"payment_params"` + // Keep the URL in top-level JSON so clients do not need to parse the protobuf just to route payment. + Orchestrator string `json:"orchestrator"` + ManifestID string `json:"manifest_id"` +} + +type liveRunnerTrickleChannelRequest struct { + Name string `json:"name"` + MimeType string `json:"mime_type"` +} + +type liveRunnerTrickleChannelsRequest struct { + Channels []liveRunnerTrickleChannelRequest `json:"channels"` +} + +type liveRunnerTrickleChannelsResponse struct { + Channels []runner.LiveRunnerTrickleChannel `json:"channels"` +} + +type deleteLiveRunnerTrickleChannelsRequest struct { + Channels []string `json:"channels"` +} + +type deleteLiveRunnerTrickleChannelsResponse struct { + Deleted []string `json:"deleted"` +} + +func (h *lphttp) ReserveLiveRunnerSession(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + priceInfo, err := manager.PaymentInfo(runnerID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + paymentRequired := priceInfo != nil + if paymentRequired && r.Header.Get(paymentHeader) == "" && r.Header.Get(segmentHeader) == "" { + h.runnerChallenge(w, r, priceInfo) + return + } + var ( + payment lpnet.Payment + segData *core.SegTranscodingMetadata + ctx = r.Context() + sessionID string + ) + if paymentRequired { + var err error + payment, segData, ctx, err = h.processPaymentAndSegmentHeaders(w, r) + if err != nil { + return + } + if string(segData.ManifestID) != segData.AuthToken.SessionId { + respondWithError(w, "mismatched manifest and auth token", http.StatusForbidden) + return + } + // for easier correlation across orch, gw, signer + sessionID = string(segData.ManifestID) + } + sessionID, _, err = manager.ReserveSession(runnerID, sessionID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + ctx = clog.AddVal(ctx, "runner_id", runnerID) + ctx = clog.AddVal(ctx, "session_id", sessionID) + if paymentRequired { + if err := h.orchestrator.ProcessPayment(ctx, payment, segData.ManifestID); err != nil { + if releaseErr := manager.ReleaseSession(runnerID, sessionID); releaseErr != nil { + clog.Errorf(ctx, "Error releasing live runner session after payment failure err=%v", releaseErr) + } + respondWithError(w, err.Error(), http.StatusBadRequest) + return + } + monitorCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + paymentReceiver := livePaymentReceiver{orchestrator: h.orchestrator} + accountPaymentFunc := func(inPixels int64) error { + err := paymentReceiver.AccountPayment(monitorCtx, &SegmentInfoReceiver{ + sender: getPaymentSender(payment), + inPixels: inPixels, + priceInfo: payment.GetExpectedPrice(), + sessionID: string(segData.ManifestID), + }) + if err != nil { + clog.Errorf(monitorCtx, "Error accounting live runner payment, releasing session err=%v", err) + if releaseErr := manager.ReleaseSession(runnerID, sessionID); releaseErr != nil { + clog.Errorf(monitorCtx, "Error releasing live runner session after payment failure err=%v", releaseErr) + } + // Stop both the ticker loop below and the LivePaymentProcessor goroutine. + cancel() + } + return err + } + paymentProcessor := NewLivePaymentProcessor(monitorCtx, h.node.LivePaymentInterval, accountPaymentFunc) + go func() { + ticker := time.NewTicker(h.node.LivePaymentInterval) + defer ticker.Stop() + defer cancel() + for { + select { + case <-ticker.C: + // Stop monitoring once the live runner session has been released + // by an explicit stop, runner cleanup, expiry, or payment failure. + if _, err := manager.RunnerEndpointForSession(runnerID, sessionID); err != nil { + return + } + paymentProcessor.process(monitorCtx) + case <-monitorCtx.Done(): + return + } + } + }() + } + controlURL := h.orchestrator.ServiceURI().JoinPath("apps", runnerID, "session", sessionID).String() + appURL := h.orchestrator.ServiceURI().JoinPath("apps", runnerID, "session", sessionID, "app").String() + data, err := json.Marshal(liveRunnerSessionResponse{SessionID: sessionID, AppURL: appURL, ControlURL: controlURL}) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + respondJsonOk(w, data) +} + +func (h *lphttp) runnerChallenge(w http.ResponseWriter, r *http.Request, priceInfo *runner.LiveRunnerPriceInfo) { + sender, err := h.runnerSender(r) + if err != nil { + respondJsonError(r.Context(), w, err, http.StatusPaymentRequired) + return + } + oInfo, err := h.runnerOrchInfo(sender, priceInfo) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + data, err := marshalLivePaymentChallengeResponse(oInfo) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusPaymentRequired) + _, _ = w.Write(data) +} + +func (h *lphttp) scopePaymentChallenge(w http.ResponseWriter, r *http.Request) (bool, error) { + sender, err := h.runnerSender(r) + if err != nil { + respondJsonError(r.Context(), w, err, http.StatusPaymentRequired) + return true, nil + } + constraints := core.NewCapabilities(nil, nil) + if h.node != nil && h.node.Capabilities != nil { + constraints = h.node.Capabilities + } + caps := newAICapabilities(core.Capability_LiveVideoToVideo, "scope", true, constraints) + oInfo, err := orchestratorInfoWithCaps(h.orchestrator, sender, h.orchestrator.ServiceURI().String(), "", caps.ToNetCapabilities()) + if err != nil { + return false, err + } + data, err := marshalLivePaymentChallengeResponse(oInfo) + if err != nil { + return false, err + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusPaymentRequired) + _, _ = w.Write(data) + return true, nil +} + +func marshalLivePaymentChallengeResponse(oInfo *lpnet.OrchestratorInfo) ([]byte, error) { + buf, err := proto.Marshal(oInfo) + if err != nil { + return nil, err + } + return json.Marshal(liveRunnerPaymentChallengeResponse{ + PaymentParams: base64.StdEncoding.EncodeToString(buf), + Orchestrator: oInfo.GetTranscoder(), + ManifestID: oInfo.GetAuthToken().GetSessionId(), + }) +} + +func (h *lphttp) processScopePayment(ctx context.Context, w http.ResponseWriter, r *http.Request) (ethcommon.Address, *lpnet.PriceInfo, string, bool, *runner.RunnerError) { + payment, err := getPayment(r.Header.Get(paymentHeader)) + if err != nil || r.Header.Get(paymentHeader) == "" { + handled, challengeErr := h.scopePaymentChallenge(w, r) + if handled { + return ethcommon.Address{}, nil, "", true, nil + } + if challengeErr != nil { + return ethcommon.Address{}, nil, "", false, &runner.RunnerError{Message: challengeErr.Error(), StatusCode: http.StatusInternalServerError} + } + return ethcommon.Address{}, nil, "", false, &runner.RunnerError{Message: err.Error(), StatusCode: http.StatusPaymentRequired} + } + sender := getPaymentSender(payment) + segData, ctx, err := verifySegCreds(ctx, h.orchestrator, r.Header.Get(segmentHeader), sender) + if err != nil { + return ethcommon.Address{}, nil, "", false, &runner.RunnerError{Message: err.Error(), StatusCode: http.StatusForbidden} + } + if string(segData.ManifestID) != segData.AuthToken.SessionId { + clog.Info(ctx, "Legacy Scope payment manifest mismatch", "manifest_id", segData.ManifestID, "auth_session_id", segData.AuthToken.SessionId) + } + manifestID := string(segData.ManifestID) + if err := h.orchestrator.ProcessPayment(ctx, payment, segData.ManifestID); err != nil { + return ethcommon.Address{}, nil, "", false, &runner.RunnerError{Message: err.Error(), StatusCode: http.StatusBadRequest} + } + if payment.GetExpectedPrice().GetPricePerUnit() > 0 && !h.orchestrator.SufficientBalance(sender, segData.ManifestID) { + return ethcommon.Address{}, nil, "", false, &runner.RunnerError{Message: "Insufficient balance", StatusCode: http.StatusBadRequest} + } + return sender, payment.GetExpectedPrice(), manifestID, false, nil +} + +func (h *lphttp) runnerSender(r *http.Request) (ethcommon.Address, error) { + addr := r.Header.Get(liveRunnerSenderHeader) + if !ethcommon.IsHexAddress(addr) { + return ethcommon.Address{}, fmt.Errorf("invalid live runner payment signer address") + } + return ethcommon.HexToAddress(addr), nil +} + +func (h *lphttp) runnerOrchInfo(sender ethcommon.Address, priceInfo *runner.LiveRunnerPriceInfo) (*lpnet.OrchestratorInfo, error) { + if priceInfo == nil { + return nil, fmt.Errorf("missing live runner price info") + } + netPriceInfo := &lpnet.PriceInfo{ + PricePerUnit: priceInfo.PricePerUnit, + PixelsPerUnit: priceInfo.PixelsPerUnit, + } + params, err := h.orchestrator.TicketParams(sender, netPriceInfo) + if err != nil { + return nil, err + } + manifestID := string(core.RandomManifestID()) + expiration := time.Now().Add(authTokenValidPeriod).Unix() + authToken := h.orchestrator.AuthToken(manifestID, expiration) + return &lpnet.OrchestratorInfo{ + Transcoder: h.orchestrator.ServiceURI().String(), + TicketParams: params, + PriceInfo: netPriceInfo, + Address: h.orchestrator.Address().Bytes(), + AuthToken: authToken, + }, nil +} + +func (h *lphttp) StopLiveRunnerSession(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + if err := manager.ReleaseSession(r.PathValue("runner_id"), r.PathValue("session_id")); err != nil { + respondWithLiveRunnerError(w, err) + return + } + w.WriteHeader(http.StatusNoContent) +} + +// PaymentForLiveRunnerSession receives payment for a live runner session and +// rejects payment once the session is no longer live. +func (h *lphttp) PaymentForLiveRunnerSession(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + sessionID := r.PathValue("session_id") + + if _, err := manager.RunnerEndpointForSession(runnerID, sessionID); err != nil { + respondWithLiveRunnerError(w, err) + return + } + + payment, segData, ctx, err := h.processPaymentAndSegmentHeaders(w, r) + if err != nil { + return + } + if string(segData.ManifestID) != sessionID { + respondWithError(w, "mismatched session and payment manifest", http.StatusForbidden) + return + } + + var netCaps *lpnet.Capabilities + if segData.Caps != nil { + netCaps = segData.Caps.ToNetCapabilities() + } + oInfo, err := orchestratorInfoWithCaps(h.orchestrator, getPaymentSender(payment), h.orchestrator.ServiceURI().String(), core.ManifestID(segData.AuthToken.SessionId), netCaps) + if err != nil { + clog.Errorf(ctx, "Error updating orchestrator info - err=%q", err) + http.Error(w, "Internal Server Error", http.StatusInternalServerError) + return + } + oInfo.AuthToken = segData.AuthToken + + if err := h.orchestrator.ProcessPayment(ctx, payment, segData.ManifestID); err != nil { + clog.Errorf(ctx, "error processing payment: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + buf, err := proto.Marshal(&lpnet.PaymentResult{Info: oInfo}) + if err != nil { + clog.Errorf(ctx, "Unable to marshal payment result err=%q", err) + return + } + clog.V(common.DEBUG).Infof(ctx, "Live runner session payment processed, current balance=%s", currentBalanceLog(h, payment, segData)) + + w.Write(buf) +} + +func (h *lphttp) StopLiveRunnerSessionInternal(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + sessionID := r.PathValue("session_id") + if err := manager.ValidSessionToken(runnerID, sessionID, r.Header.Get("Livepeer-Session-Token")); err != nil { + respondWithLiveRunnerError(w, err) + return + } + if err := manager.ReleaseSession(runnerID, sessionID); err != nil { + respondWithLiveRunnerError(w, err) + return + } + w.WriteHeader(http.StatusNoContent) +} + +func (h *lphttp) CreateLiveRunnerTrickleChannel(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + sessionID := r.PathValue("session_id") + if err := manager.ValidSessionToken(runnerID, sessionID, r.Header.Get("Livepeer-Session-Token")); err != nil { + respondWithLiveRunnerError(w, err) + return + } + + var req liveRunnerTrickleChannelsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + respondWithError(w, err.Error(), http.StatusBadRequest) + return + } + if len(req.Channels) == 0 { + respondWithError(w, "channels is required", http.StatusBadRequest) + return + } + if len(req.Channels) > maxLiveRunnerTrickleChannelsPerRequest { + respondWithError(w, fmt.Sprintf("channels cannot contain more than %d entries", maxLiveRunnerTrickleChannelsPerRequest), http.StatusBadRequest) + return + } + + channels := make([]runner.LiveRunnerTrickleChannel, 0, len(req.Channels)) + for _, channelReq := range req.Channels { + channel, err := manager.CreateTrickleChannel( + runnerID, + sessionID, + channelReq.Name, + channelReq.MimeType, + ) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + channels = append(channels, channel) + } + data, err := json.Marshal(liveRunnerTrickleChannelsResponse{Channels: channels}) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + respondJsonOk(w, data) +} + +func (h *lphttp) DeleteLiveRunnerTrickleChannels(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + runnerID := r.PathValue("runner_id") + sessionID := r.PathValue("session_id") + if err := manager.ValidSessionToken(runnerID, sessionID, r.Header.Get("Livepeer-Session-Token")); err != nil { + respondWithLiveRunnerError(w, err) + return + } + + var req deleteLiveRunnerTrickleChannelsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + respondWithError(w, err.Error(), http.StatusBadRequest) + return + } + if len(req.Channels) == 0 { + respondWithError(w, "channels is required", http.StatusBadRequest) + return + } + if len(req.Channels) > maxLiveRunnerTrickleChannelsPerRequest { + respondWithError(w, fmt.Sprintf("channels cannot contain more than %d entries", maxLiveRunnerTrickleChannelsPerRequest), http.StatusBadRequest) + return + } + for _, channelName := range req.Channels { + if err := manager.DeleteTrickleChannel(runnerID, sessionID, channelName); err != nil { + respondWithLiveRunnerError(w, err) + return + } + } + data, err := json.Marshal(deleteLiveRunnerTrickleChannelsResponse{Deleted: req.Channels}) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + respondJsonOk(w, data) +} + +func (h *lphttp) ProxyLiveRunnerSession(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + + runnerID := r.PathValue("runner_id") + sessionID := r.PathValue("session_id") + endpoint, err := manager.RunnerEndpointForSession(runnerID, sessionID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + sessionToken, err := manager.SessionTokenForSession(runnerID, sessionID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + + h.proxyLiveRunner(w, r, runnerID, sessionID, sessionToken, endpoint) +} + +func (h *lphttp) ProxyLiveRunnerSingleShot(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + if !ok { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + + runnerID := r.PathValue("runner_id") + mode, err := manager.RunnerMode(runnerID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + if mode != runner.LiveRunnerModeSingleShot { + respondWithError(w, "runner is not single-shot", http.StatusBadRequest) + return + } + + sessionID, endpoint, err := manager.ReserveSession(runnerID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + defer func() { + if err := manager.ReleaseSession(runnerID, sessionID); err != nil { + slog.Error("error releasing single-shot session", "runner_id", runnerID, "session_id", sessionID, "err", err) + } + }() + + sessionToken, err := manager.SessionTokenForSession(runnerID, sessionID) + if err != nil { + respondWithLiveRunnerError(w, err) + return + } + h.proxyLiveRunner(w, r, runnerID, sessionID, sessionToken, endpoint) +} + +func (h *lphttp) proxyLiveRunner(w http.ResponseWriter, r *http.Request, runnerID, sessionID, sessionToken, endpoint string) { + target, err := url2.Parse(endpoint) + if err != nil { + respondWithError(w, err.Error(), http.StatusBadGateway) + return + } + appPath := r.PathValue("app_path") + proxyPath, err := url2.JoinPath("/", target.Path, appPath) + if err != nil { + respondWithError(w, err.Error(), http.StatusBadGateway) + return + } + sessionControlURL := liveRunnerURI(h.node, h.orchestrator).JoinPath("runner", runnerID, "session", sessionID).String() + proxy := &httputil.ReverseProxy{ + Rewrite: func(req *httputil.ProxyRequest) { + req.Out.URL.Scheme = target.Scheme + req.Out.URL.Host = target.Host + req.Out.URL.Path = proxyPath + req.Out.URL.RawPath = "" + req.Out.URL.RawQuery = req.In.URL.RawQuery + req.Out.Host = target.Host + req.Out.Header.Set("Livepeer-Runner-Route", runnerID) + req.Out.Header.Set("Livepeer-Session-Id", sessionID) + req.Out.Header.Set("Livepeer-Session-Token", sessionToken) + req.Out.Header.Set("Livepeer-Session-Control", sessionControlURL) + }, + ErrorHandler: func(w http.ResponseWriter, req *http.Request, err error) { + respondWithError(w, err.Error(), http.StatusBadGateway) + }, + } + proxy.ServeHTTP(w, r) +} + +func respondWithLiveRunnerError(w http.ResponseWriter, err error) { + var runnerErr *runner.RunnerError + if errors.As(err, &runnerErr) { + respondWithError(w, runnerErr.Error(), runnerErr.StatusCode) + return + } + respondWithError(w, err.Error(), http.StatusInternalServerError) +} + +type liveRunnerDiscoveryEntry struct { + Address string `json:"address,omitempty"` + Runners []runner.LiveRunnerDiscoveryRunner `json:"runners,omitempty"` +} + +func (h *lphttp) DiscoverLiveRunners(w http.ResponseWriter, r *http.Request) { + manager, ok := h.liveRunnerManager() + var hasServerlessWorker bool + if h.node != nil { + _, hasServerlessWorker = h.node.AIWorker.(*worker.ServerlessWorker) + } + if !ok && !hasServerlessWorker { + respondWithError(w, "live runners are not supported", http.StatusNotFound) + return + } + + suri := h.orchestrator.ServiceURI() + + var runners []runner.LiveRunnerDiscoveryRunner + if ok { + runners = append(runners, manager.Runners()...) + } + if hasServerlessWorker { + pipeline := "live-video-to-video" + capability := core.Capability_LiveVideoToVideo + var capConstraints *core.CapabilityConstraints + if h.node.Capabilities != nil { + capConstraints = h.node.Capabilities.PerCapability()[capability] + } + + versionFallback := "" + versions := h.node.AIWorker.Version() + for _, version := range versions { + if versionFallback == "" && version.Version != "" { + versionFallback = version.Version + } + } + + if capConstraints != nil { + for modelID := range capConstraints.Models { + versionString := versionFallback + for _, version := range versions { + if version.Pipeline == pipeline && version.ModelId == modelID && version.Version != "" { + versionString = version.Version + break + } + } + + var priceInfo *runner.LiveRunnerPriceInfo + price := h.node.GetBasePriceForCap("default", capability, modelID) + if price != nil { + priceInt64, err := common.PriceToInt64(price) + if err != nil { + glog.Errorf("error converting discovery price for capability %v modelID=%v err=%v", core.CapabilityNameLookup[capability], modelID, err) + } else { + priceInfo = &runner.LiveRunnerPriceInfo{ + PricePerUnit: priceInt64.Num().Int64(), + PixelsPerUnit: priceInt64.Denom().Int64(), + Unit: "WEI", + } + } + } + + capacity := h.node.AIWorker.GetLiveAICapacity(pipeline, modelID) + runners = append(runners, runner.LiveRunnerDiscoveryRunner{ + URL: suri.JoinPath("scope").String(), + GPU: &runner.LiveRunnerGPU{Name: "H100"}, + App: pipeline + "/" + modelID, + Version: versionString, + Capacity: capacity.ContainersIdle + capacity.ContainersInUse, + CapacityUsed: capacity.ContainersInUse, + CapacityAvailable: capacity.ContainersIdle, + PriceInfo: priceInfo, + }) + } + } + } + + resp := []liveRunnerDiscoveryEntry{{ + Address: suri.String(), + Runners: runners, + }} + data, err := json.Marshal(resp) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + return + } + respondJsonOk(w, data) +} + // aiHttpHandle handles AI requests by decoding the request body and processing it. func aiHttpHandle[I any](h *lphttp, decoderFunc func(*I, *http.Request) error) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -94,6 +807,19 @@ func aiHttpHandle[I any](h *lphttp, decoderFunc func(*I, *http.Request) error) h }) } +func serverlessHandshakeHTTPStatus(err error) (int, bool) { + var hsErr *worker.ServerlessHandshakeError + if !errors.As(err, &hsErr) { + return 0, false + } + switch hsErr.StatusCode { + case http.StatusBadRequest, http.StatusUnauthorized: + return hsErr.StatusCode, true + default: + return 0, false + } +} + func (h *lphttp) StartLiveVideoToVideo() http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { startTime := time.Now() @@ -286,6 +1012,183 @@ func (h *lphttp) StartLiveVideoToVideo() http.Handler { }) } +func (h *lphttp) StartScope() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + startTime := time.Now() + remoteAddr := getRemoteAddr(r) + ctx := clog.AddVal(context.Background(), clog.ClientIP, remoteAddr) + + gatewayRequestID := r.Header.Get("requestID") + + var req worker.GenLiveVideoToVideoJSONRequestBody + r.Body = http.MaxBytesReader(w, r.Body, maxScopeRequestBodySize) + if err := jsonDecoder(&req, r); err != nil { + respondWithError(w, err.Error(), http.StatusBadRequest) + return + } + + if req.GatewayRequestId != nil && *req.GatewayRequestId != "" { + gatewayRequestID = *req.GatewayRequestId + } + ctx = clog.AddVal(ctx, "request_id", gatewayRequestID) + ctx = clog.AddVal(ctx, "app", "scope") + orch := h.orchestrator + pipeline := "live-video-to-video" + modelID := "scope" + + var ( + manifestID string + sender ethcommon.Address + priceInfo *lpnet.PriceInfo + ) + if h.node != nil && h.node.Eth != nil { + payer, price, paidManifestID, handled, httpErr := h.processScopePayment(ctx, w, r) + if handled { + return + } + if httpErr != nil { + respondWithError(w, httpErr.Error(), httpErr.StatusCode) + return + } + manifestID = paidManifestID + sender = payer + priceInfo = price + } else { + manifestID = string(core.RandomManifestID()) + } + + ctx = clog.AddVal(ctx, "manifest_id", manifestID) + clog.Info(ctx, "Received Scope request") + + // Create storage for the request (for AI Workers, must run before CheckAICapacity) + err := orch.CreateStorageForRequest(manifestID) + if err != nil { + respondWithError(w, "Could not create storage to receive results", http.StatusInternalServerError) + return + } + + // Check if there is capacity for the request + hasCapacity, _ := orch.CheckAICapacity(pipeline, modelID) + if !hasCapacity { + clog.Errorf(ctx, "Insufficient capacity for pipeline=%v modelID=%v", pipeline, modelID) + respondWithError(w, "insufficient capacity", http.StatusServiceUnavailable) + return + } + + // Start trickle server for scope + baseURL := orch.ServiceURI().JoinPath(TrickleHTTPPath, manifestID).String() + controlURL := baseURL + "-control" + eventsURL := baseURL + "-events" + + // Scope only pre-creates control + events channels. + controlPubCh := trickle.NewLocalPublisher(h.trickleSrv, manifestID+"-control", "application/json") + controlPubCh.CreateChannel() + eventsCh := trickle.NewLocalPublisher(h.trickleSrv, manifestID+"-events", "application/json") + eventsCh.CreateChannel() + + ctx, cancel := context.WithCancel(ctx) + closeSession := func() { + eventsCh.Close() + controlPubCh.Close() + cancel() + } + + // Start payment receiver which accounts the payments and stops the stream if the payment is insufficient + var paymentProcessor *LivePaymentProcessor + if priceInfo != nil && priceInfo.PricePerUnit != 0 { + paymentReceiver := livePaymentReceiver{orchestrator: h.orchestrator} + accountPaymentFunc := func(inPixels int64) error { + err := paymentReceiver.AccountPayment(ctx, &SegmentInfoReceiver{ + sender: sender, + inPixels: inPixels, + priceInfo: priceInfo, + sessionID: manifestID, + }) + if err != nil { + clog.Errorf(ctx, "Error accounting payment, stopping stream processing", err) + closeSession() + } + return err + } + paymentProcessor = NewLivePaymentProcessor(ctx, h.node.LivePaymentInterval, accountPaymentFunc) + } else { + clog.Warningf(ctx, "No price info found for model %v, Orchestrator will not charge for video processing", modelID) + } + + // For every event segment, check payments + go func() { + sub := trickle.NewLocalSubscriber(h.trickleSrv, manifestID+"-events") + for { + // Set seq to next segment in case the subscriber is outside + // the server's retention window + sub.SetSeq(-1) + segment, err := sub.Read() + if err != nil { + clog.Infof(ctx, "Error getting local trickle segment err=%v", err) + closeSession() + return + } + if paymentProcessor != nil { + paymentProcessor.process(ctx) + } + // read the segment so we know when it is complete, otherwise sub.Read() + // would rapidly request follow-on segments that do not yet exist + io.Copy(io.Discard, segment.Reader) + } + }() + + // Prepare request to worker + workerControlURL := overwriteHost(h.node.LiveAITrickleHostForRunner, controlURL) + workerEventsURL := overwriteHost(h.node.LiveAITrickleHostForRunner, eventsURL) + + workerReq := worker.LiveVideoToVideoParams{ + EventsUrl: &workerEventsURL, + ControlUrl: &workerControlURL, + Params: req.Params, + GatewayRequestId: &gatewayRequestID, + ManifestId: &manifestID, + } + + _, err = orch.LiveVideoToVideo(ctx, manifestID, workerReq) + if err != nil { + if monitor.Enabled { + monitor.AIProcessingError(err.Error(), pipeline, modelID, ethcommon.Address{}.String()) + } + + closeSession() + if statusCode, ok := serverlessHandshakeHTTPStatus(err); ok { + respondWithError(w, err.Error(), statusCode) + } else { + respondWithError(w, err.Error(), http.StatusInternalServerError) + } + return + } + + jsonData, err := json.Marshal(&worker.LiveVideoToVideoResponse{ + ControlUrl: &controlURL, + EventsUrl: &eventsURL, + ManifestId: &manifestID, + }) + if err != nil { + respondWithError(w, err.Error(), http.StatusInternalServerError) + closeSession() + return + } + + took := time.Since(startTime) + clog.Info(ctx, "Processed Scope request", "took", took) + respondJsonOk(w, jsonData) + }) +} + +func liveRunnerURI(node *core.LivepeerNode, orch Orchestrator) *url2.URL { + if node != nil && node.LiveRunnerAddr != nil { + v := *node.LiveRunnerAddr + return &v + } + return orch.ServiceURI() +} + // overwriteHost is used to overwrite the trickle host, because it may be different for runner // runner may run inside Docker container, in a different network, or even on a different machine func overwriteHost(hostOverwrite, url string) string { diff --git a/server/ai_http_test.go b/server/ai_http_test.go index 4a3bb66a26..07b91704ea 100644 --- a/server/ai_http_test.go +++ b/server/ai_http_test.go @@ -1,19 +1,71 @@ package server import ( + "bytes" + "context" "crypto/tls" + "encoding/base64" + "encoding/json" + "errors" + "fmt" "io" + "math/big" "net/http" "net/http/httptest" "net/url" + "strings" "testing" + "testing/synctest" "time" + ethcommon "github.com/ethereum/go-ethereum/common" + "github.com/golang/protobuf/proto" + "github.com/livepeer/go-livepeer/ai/runner" + "github.com/livepeer/go-livepeer/ai/worker" "github.com/livepeer/go-livepeer/core" + "github.com/livepeer/go-livepeer/eth" + lpnet "github.com/livepeer/go-livepeer/net" + "github.com/livepeer/go-livepeer/trickle" + "github.com/livepeer/go-tools/drivers" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +type stubPriceFeedWatcher struct { + price eth.PriceData +} + +type testLiveRunnerHost struct { + orchestrator Orchestrator + liveRunnerURI *url.URL +} + +func (h testLiveRunnerHost) ServiceURI() *url.URL { + return h.orchestrator.ServiceURI() +} + +func (h testLiveRunnerHost) LiveRunnerURI() *url.URL { + if h.liveRunnerURI != nil { + v := *h.liveRunnerURI + return &v + } + return h.ServiceURI() +} + +func (h testLiveRunnerHost) RegistrationSecret() string { + return h.orchestrator.RegistrationSecret() +} + +func (s stubPriceFeedWatcher) Currencies() (string, string, error) { + return "ETH", "USD", nil +} + +func (s stubPriceFeedWatcher) Current() (eth.PriceData, error) { + return s.price, nil +} + +func (s stubPriceFeedWatcher) Subscribe(context.Context, chan<- eth.PriceData) {} + func TestAIWorkerResults_ErrorsWhenAuthHeaderMissing(t *testing.T) { var l lphttp @@ -123,3 +175,1724 @@ func TestAIWorkerResults_BadRequestType(t *testing.T) { assert.Equal(protoVerAIWorker, headers.Get("Authorization")) assert.Equal("AI request validation failed for", string(body)[0:32]) } + +func TestLiveRunnerDiscoveryEndpoint(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + liveRunnerAddr, err := url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + lp.node.LiveRunnerAddr = liveRunnerAddr + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + _, err = manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + Version: "1.2.3", + Status: "ready", + GPU: &runner.LiveRunnerGPU{Name: "NVIDIA L40S", VRAMMB: 46068}, + App: "live-video-to-video/scope", + Capacity: 2, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "USD", + }, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Len(t, resp[0].Runners, 1) + require.Equal(t, "http://localhost:1234/apps/"+t.Name()+"/session", resp[0].Runners[0].URL) + require.Equal(t, "live-video-to-video/scope", resp[0].Runners[0].App) + require.Equal(t, 2, resp[0].Runners[0].Capacity) + require.Equal(t, 0, resp[0].Runners[0].CapacityUsed) + require.Equal(t, 2, resp[0].Runners[0].CapacityAvailable) + require.NotContains(t, w.Body.String(), "endpoint") + require.NotContains(t, w.Body.String(), "price_info") +} + +func TestLiveRunnerDiscoveryReportsCapacityUsed(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + _, err := manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + Status: "ready", + App: "live-video-to-video/scope", + Capacity: 2, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) + _, _, err = manager.ReserveSession(t.Name()) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Len(t, resp[0].Runners, 1) + require.Equal(t, 2, resp[0].Runners[0].Capacity) + require.Equal(t, 1, resp[0].Runners[0].CapacityUsed) + require.Equal(t, 1, resp[0].Runners[0].CapacityAvailable) + require.Contains(t, w.Body.String(), `"capacity":2`) + require.Contains(t, w.Body.String(), `"capacity_used":1`) + require.Contains(t, w.Body.String(), `"capacity_available":1`) +} + +func TestLiveRunnerDiscoverySingleShotReturnsProxiedURL(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + _, err := manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + Status: "ready", + Mode: runner.LiveRunnerModeSingleShot, + App: "live-video-to-video/scope", + Capacity: 1, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Len(t, resp[0].Runners, 1) + require.Equal(t, "http://localhost:1234/apps/"+t.Name()+"/app", resp[0].Runners[0].URL) + require.Equal(t, runner.LiveRunnerModeSingleShot, resp[0].Runners[0].Mode) + require.Equal(t, 1, resp[0].Runners[0].Capacity) + require.Equal(t, 0, resp[0].Runners[0].CapacityUsed) + require.Equal(t, 1, resp[0].Runners[0].CapacityAvailable) + require.NotContains(t, w.Body.String(), "endpoint") + require.NotContains(t, w.Body.String(), "price_info") +} + +func TestLiveRunnerDiscoveryOnchainIncludesPriceInfo(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}, Onchain: true}) + t.Cleanup(manager.Stop) + lp.node.LiveRunnerManager = manager + trickleBaseURL := lp.orchestrator.ServiceURI().JoinPath(TrickleHTTPPath).String() + manager.SetTrickleServer(lp.trickleSrv, trickleBaseURL, trickleBaseURL) + prevWatcher := core.PriceFeedWatcher + core.PriceFeedWatcher = stubPriceFeedWatcher{price: eth.PriceData{Price: big.NewRat(2000, 1)}} + defer func() { core.PriceFeedWatcher = prevWatcher }() + _, err := manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + Version: "1.2.3", + Status: "ready", + App: "live-video-to-video/scope", + Capacity: 2, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "USD", + }, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Len(t, resp[0].Runners, 1) + require.Equal(t, "http://localhost:1234/apps/"+t.Name()+"/session", resp[0].Runners[0].URL) + require.Equal(t, 2, resp[0].Runners[0].Capacity) + require.Equal(t, 0, resp[0].Runners[0].CapacityUsed) + require.Equal(t, 2, resp[0].Runners[0].CapacityAvailable) + require.NotNil(t, resp[0].Runners[0].PriceInfo) + require.NotZero(t, resp[0].Runners[0].PriceInfo.PricePerUnit) + require.Contains(t, w.Body.String(), "price_info") +} + +func TestLiveRunnerDiscoveryServerlessWorker(t *testing.T) { + lp := newServerlessLiveRunnerHTTP(t, false, 3) + lp.node.SetBasePriceForCap("default", core.Capability_LiveVideoToVideo, "scope", core.NewFixedPrice(big.NewRat(7, 1))) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Equal(t, "http://localhost:1234", resp[0].Address) + require.Len(t, resp[0].Runners, 1) + + discoveryRunner := resp[0].Runners[0] + require.Equal(t, "http://localhost:1234/scope", discoveryRunner.URL) + require.NotNil(t, discoveryRunner.GPU) + require.Equal(t, "H100", discoveryRunner.GPU.Name) + require.Equal(t, "live-video-to-video/scope", discoveryRunner.App) + require.Equal(t, "serverless-1.0.0", discoveryRunner.Version) + require.Equal(t, 3, discoveryRunner.Capacity) + require.Equal(t, 0, discoveryRunner.CapacityUsed) + require.Equal(t, 3, discoveryRunner.CapacityAvailable) + require.NotNil(t, discoveryRunner.PriceInfo) + require.Equal(t, int64(7), discoveryRunner.PriceInfo.PricePerUnit) + require.Equal(t, int64(1), discoveryRunner.PriceInfo.PixelsPerUnit) +} + +func TestLiveRunnerDiscoveryReturnsHeartbeatAndServerlessRunners(t *testing.T) { + lp := newServerlessLiveRunnerHTTP(t, true, 2) + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + prevWatcher := core.PriceFeedWatcher + core.PriceFeedWatcher = stubPriceFeedWatcher{price: eth.PriceData{Price: big.NewRat(2000, 1)}} + defer func() { core.PriceFeedWatcher = prevWatcher }() + _, err := manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + Version: "1.2.3", + Status: "ready", + GPU: &runner.LiveRunnerGPU{Name: "NVIDIA L40S", VRAMMB: 46068}, + App: "live-video-to-video/scope", + Capacity: 1, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "USD", + }, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp []liveRunnerDiscoveryEntry + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp, 1) + require.Len(t, resp[0].Runners, 2) + require.Equal(t, "http://localhost:1234/apps/"+t.Name()+"/session", resp[0].Runners[0].URL) + require.Equal(t, 1, resp[0].Runners[0].Capacity) + require.Equal(t, 0, resp[0].Runners[0].CapacityUsed) + require.Equal(t, 1, resp[0].Runners[0].CapacityAvailable) + require.Equal(t, "http://localhost:1234/scope", resp[0].Runners[1].URL) + require.Equal(t, "H100", resp[0].Runners[1].GPU.Name) + require.Equal(t, 2, resp[0].Runners[1].Capacity) + require.Equal(t, 0, resp[0].Runners[1].CapacityUsed) + require.Equal(t, 2, resp[0].Runners[1].CapacityAvailable) +} + +func TestLiveRunnerHeartbeat(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + body, err := json.Marshal(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 3, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "USD", + }, + SessionIDs: []string{"runner-reported-session"}, + }) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader(body)) + req.Header.Set("Authorization", lp.orchestrator.RegistrationSecret()) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp runner.LiveRunnerHeartbeatResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, t.Name(), resp.RunnerID) + require.Equal(t, lp.orchestrator.ServiceURI().String(), resp.Orchestrator) + require.NotEmpty(t, resp.HeartbeatSecret) + require.Empty(t, resp.SessionIDs) + require.NotNil(t, resp.O2R) + require.Equal(t, runner.LiveRunnerTrickleOrchestratorToRunner, resp.O2R.Name) + for _, channel := range []runner.LiveRunnerTrickleChannel{*resp.O2R} { + require.True(t, strings.HasPrefix(channel.ChannelName, resp.RunnerID+"-")) + require.True(t, strings.HasSuffix(channel.ChannelName, "-"+channel.Name)) + require.NotEqual(t, resp.RunnerID+"-"+channel.Name, channel.ChannelName) + require.Equal(t, "application/json", channel.MimeType) + require.Equal(t, "http://localhost:1234/ai/trickle/"+channel.ChannelName, channel.URL) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/ai/trickle/"+channel.ChannelName+"/next", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + } + + // Check missing auth after bootstrap. + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader(body)) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + _, _, err = manager.ReserveSession(t.Name(), "session-z-oldest") + require.NoError(t, err) + time.Sleep(time.Millisecond) + _, _, err = manager.ReserveSession(t.Name(), "session-a-newest") + require.NoError(t, err) + + // Check follow-up heartbeat auth. + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader(body)) + req.Header.Set("Authorization", resp.HeartbeatSecret) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var nextResp runner.LiveRunnerHeartbeatResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &nextResp)) + require.Equal(t, t.Name(), nextResp.RunnerID) + require.Empty(t, nextResp.HeartbeatSecret) + require.Nil(t, nextResp.O2R) + require.Equal(t, []string{"session-z-oldest", "session-a-newest"}, nextResp.SessionIDs) +} + +func TestLiveRunnerHeartbeatUsesLiveRunnerAddr(t *testing.T) { + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + node.LiveRunnerAddr, err = url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + orch := newStubOrchestrator() + orch.serviceURI = "https://public.example.com" + orch.secret = "live-runner-secret" + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{ + Host: testLiveRunnerHost{orchestrator: lp.orchestrator, liveRunnerURI: node.LiveRunnerAddr}, + }) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + require.NoError(t, startAIServer(lp)) + + body, err := json.Marshal(runner.LiveRunnerHeartbeatRequest{ + RunnerID: t.Name(), + RunnerURL: "https://runner.example.com", + App: "live-video-to-video/scope", + Capacity: 1, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "WEI", + }, + }) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader(body)) + req.Header.Set("Authorization", lp.orchestrator.RegistrationSecret()) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp runner.LiveRunnerHeartbeatResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, "http://go-livepeer:8935", resp.Orchestrator) +} + +func TestLiveRunnerHeartbeatRejectsMissingPriceInfoUnit(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}, Onchain: true}) + t.Cleanup(manager.Stop) + lp.node.LiveRunnerManager = manager + body := []byte(fmt.Sprintf(`{ + "runner_id":%q, + "runner_url":"https://runner.example.com", + "price_info":{"price_per_unit":1,"pixels_per_unit":1}, + "app":"live-video-to-video/scope", + "capacity":1, + "pricing":{"usd_per_hour":10} + }`, t.Name())) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader(body)) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + require.Contains(t, w.Body.String(), "price_info.unit") +} + +func TestLiveRunnerEndpointsUnsupportedWithoutManager(t *testing.T) { + lp := newLiveRunnerHTTP(t, false) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "live runners are not supported") + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader([]byte(`{}`))) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "live runners are not supported") +} + +func TestLiveRunnerDiscoverySupportedWithoutManagerWhenServerlessWorkerPresent(t *testing.T) { + lp := newServerlessLiveRunnerHTTP(t, false, 1) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/discovery", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/runners/heartbeat", bytes.NewReader([]byte(`{}`))) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "live runners are not supported") +} + +func TestLiveRunnerReserveSession(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerSessionResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotEmpty(t, resp.SessionID) + require.Equal(t, "http://localhost:1234/apps/runner-1/session/"+resp.SessionID+"/app", resp.AppURL) + require.Equal(t, "http://localhost:1234/apps/runner-1/session/"+resp.SessionID, resp.ControlURL) +} + +func TestLiveRunnerReserveSessionOnchainReturnsPaymentChallenge(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, liveRunnerSenderHeaders(lp.orchestrator.(*stubOrchestrator))) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusPaymentRequired, w.Code) + challenge, oInfo := decodeLiveRunnerPaymentChallenge(t, w.Body.Bytes()) + require.NotEmpty(t, challenge.PaymentParams) + require.Equal(t, lp.orchestrator.ServiceURI().String(), challenge.Orchestrator) + require.NotEmpty(t, challenge.ManifestID) + require.Equal(t, challenge.ManifestID, oInfo.GetAuthToken().GetSessionId()) + require.NotNil(t, oInfo.GetTicketParams()) + require.NotNil(t, oInfo.GetPriceInfo()) + require.Equal(t, int64(10), oInfo.GetPriceInfo().GetPricePerUnit()) + require.Equal(t, int64(1), oInfo.GetPriceInfo().GetPixelsPerUnit()) + require.Equal(t, challenge.Orchestrator, oInfo.GetTranscoder()) + require.Nil(t, oInfo.GetCapabilities()) +} + +func TestLiveRunnerReserveSessionOnchainAcceptsPaidReservation(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + challenge, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + orch := lp.orchestrator.(*stubOrchestrator) + headers := liveRunnerReservationPaymentHeaders(t, orch, oInfo.GetAuthToken(), challenge.ManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerSessionResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, challenge.ManifestID, resp.SessionID) + require.Equal(t, "http://localhost:1234/apps/runner-1/session/"+resp.SessionID+"/app", resp.AppURL) + require.Equal(t, "http://localhost:1234/apps/runner-1/session/"+resp.SessionID, resp.ControlURL) +} + +func TestLiveRunnerSessionPaymentAcceptsPayment(t *testing.T) { + oldStorage := drivers.NodeStorage + drivers.NodeStorage = drivers.NewMemoryDriver(nil) + t.Cleanup(func() { drivers.NodeStorage = oldStorage }) + + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(100, 1) + + challenge, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + headers := liveRunnerReservationPaymentHeaders(t, orch, oInfo.GetAuthToken(), challenge.ManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+challenge.ManifestID+"/payment", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var paymentResult lpnet.PaymentResult + require.NoError(t, proto.Unmarshal(w.Body.Bytes(), &paymentResult)) + require.NotNil(t, paymentResult.GetInfo()) + require.Equal(t, challenge.ManifestID, paymentResult.GetInfo().GetAuthToken().GetSessionId()) + + balance := orch.Balance(orch.Address(), core.ManifestID(challenge.ManifestID)) + require.NotNil(t, balance) + require.Equal(t, "200", balance.FloatString(0)) +} + +func TestLiveRunnerSessionPaymentRejectsMissingSession(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + challenge, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + headers := liveRunnerReservationPaymentHeaders(t, lp.orchestrator.(*stubOrchestrator), oInfo.GetAuthToken(), challenge.ManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/missing/payment", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "runner session not found") +} + +func TestLiveRunnerSessionPaymentRejectsReleasedSession(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + challenge, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + headers := liveRunnerReservationPaymentHeaders(t, lp.orchestrator.(*stubOrchestrator), oInfo.GetAuthToken(), challenge.ManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+challenge.ManifestID+"/stop", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNoContent, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+challenge.ManifestID+"/payment", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "runner session not found") +} + +func TestLiveRunnerSessionPaymentRejectsManifestMismatch(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{Capacity: 2}) + + sessionID := reservePaidLiveRunnerSession(t, lp, "runner-1", &lpnet.PriceInfo{PricePerUnit: 10, PixelsPerUnit: 1}) + + _, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + headers := liveRunnerReservationPaymentHeaders(t, lp.orchestrator.(*stubOrchestrator), oInfo.GetAuthToken(), "different-manifest") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/payment", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusForbidden, w.Code) + require.Contains(t, w.Body.String(), "mismatched session and payment manifest") +} + +func TestLiveRunnerSessionPaymentRejectsBadPaymentHeader(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + sessionID := reservePaidLiveRunnerSession(t, lp, "runner-1", &lpnet.PriceInfo{PricePerUnit: 10, PixelsPerUnit: 1}) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/payment", nil) + req.Header.Set(paymentHeader, "not-base64") + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusPaymentRequired, w.Code) + require.Contains(t, w.Body.String(), "base64 decode error") +} + +func TestLiveRunnerPaidSessionMonitorDebitsBalance(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + lp.node.LivePaymentInterval = time.Second + registerLiveRunnerForSession(t, lp, nil) + manager := lp.node.LiveRunnerManager.(*runner.LiveRunnerRegistry) + defer manager.Stop() + + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(3, 1) + priceInfo := liveRunnerTestPricePerSecond(1) + + sessionID := reservePaidLiveRunnerSession(t, lp, "runner-1", priceInfo) + balance := orch.Balance(orch.Address(), core.ManifestID(sessionID)) + require.NotNil(t, balance) + require.Equal(t, "3", balance.FloatString(0)) + + time.Sleep(time.Second) + synctest.Wait() + + balance = orch.Balance(orch.Address(), core.ManifestID(sessionID)) + require.NotNil(t, balance) + require.Equal(t, "2", balance.FloatString(0)) + + require.NoError(t, manager.ReleaseSession("runner-1", sessionID)) + time.Sleep(time.Second) + synctest.Wait() + }) +} + +func TestLiveRunnerPaidSessionMonitorReleasesOnInsufficientBalance(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + lp.node.LivePaymentInterval = time.Second + registerLiveRunnerForSession(t, lp, nil) + manager := lp.node.LiveRunnerManager.(*runner.LiveRunnerRegistry) + defer manager.Stop() + + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(0, 1) + + sessionID := reservePaidLiveRunnerSession(t, lp, "runner-1", liveRunnerTestPricePerSecond(1)) + + time.Sleep(time.Second) + synctest.Wait() + + _, err := manager.RunnerEndpointForSession("runner-1", sessionID) + var runnerErr *runner.RunnerError + require.True(t, errors.As(err, &runnerErr), "expected runner error, got %v", err) + require.Equal(t, http.StatusNotFound, runnerErr.StatusCode) + }) +} + +func TestLiveRunnerPaidSessionMonitorExitsAfterManualStop(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + lp.node.LivePaymentInterval = time.Second + registerLiveRunnerForSession(t, lp, nil) + manager := lp.node.LiveRunnerManager.(*runner.LiveRunnerRegistry) + defer manager.Stop() + + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(3, 1) + + sessionID := reservePaidLiveRunnerSession(t, lp, "runner-1", liveRunnerTestPricePerSecond(1)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/stop", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNoContent, w.Code) + + time.Sleep(time.Second) + synctest.Wait() + + balance := orch.Balance(orch.Address(), core.ManifestID(sessionID)) + require.NotNil(t, balance) + require.Equal(t, "3", balance.FloatString(0)) + }) +} + +func TestLiveRunnerOffchainSessionDoesNotStartPaymentMonitor(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + lp.node.LivePaymentInterval = time.Second + registerLiveRunnerForSession(t, lp, nil) + manager := lp.node.LiveRunnerManager.(*runner.LiveRunnerRegistry) + defer manager.Stop() + + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + time.Sleep(time.Second) + synctest.Wait() + + if _, err := manager.RunnerEndpointForSession("runner-1", sessionID); err != nil { + t.Fatal(err) + } + require.NoError(t, manager.ReleaseSession("runner-1", sessionID)) + }) +} + +func TestLiveRunnerReserveSessionRejectsManifestAuthMismatch(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + registerLiveRunnerForSession(t, lp, nil) + + _, oInfo := requestLiveRunnerPaymentChallenge(t, lp, "runner-1") + orch := lp.orchestrator.(*stubOrchestrator) + headers := liveRunnerReservationPaymentHeaders(t, orch, oInfo.GetAuthToken(), "different-manifest") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusForbidden, w.Code) + require.Contains(t, w.Body.String(), "mismatched manifest and auth token") +} + +func TestScopePaymentChallenge(t *testing.T) { + lp := newScopeHTTP(t) + + challenge, oInfo := requestScopePaymentChallenge(t, lp) + + require.NotEmpty(t, challenge.PaymentParams) + require.Equal(t, lp.orchestrator.ServiceURI().String(), challenge.Orchestrator) + require.NotEmpty(t, challenge.ManifestID) + require.Equal(t, challenge.ManifestID, oInfo.GetAuthToken().GetSessionId()) + require.NotNil(t, oInfo.GetTicketParams()) + require.NotNil(t, oInfo.GetPriceInfo()) + require.Equal(t, int64(4), oInfo.GetPriceInfo().GetPricePerUnit()) + require.Equal(t, int64(1), oInfo.GetPriceInfo().GetPixelsPerUnit()) +} + +func TestLiveRunnerReserveSessionOnchainUsesPublicServiceURIForPaymentChallenge(t *testing.T) { + lp := newLiveRunnerHTTPOnchain(t) + liveRunnerAddr, err := url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + lp.node.LiveRunnerAddr = liveRunnerAddr + lp.orchestrator.(*stubOrchestrator).serviceURI = "https://public.example.com" + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + setRequestHeaders(req, liveRunnerSenderHeaders(lp.orchestrator.(*stubOrchestrator))) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusPaymentRequired, w.Code) + challenge, oInfo := decodeLiveRunnerPaymentChallenge(t, w.Body.Bytes()) + require.Equal(t, "https://public.example.com", challenge.Orchestrator) + require.Equal(t, challenge.Orchestrator, oInfo.GetTranscoder()) +} + +func TestScopePaymentChallengeUsesPublicServiceURI(t *testing.T) { + lp := newScopeHTTP(t) + liveRunnerAddr, err := url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + lp.node.LiveRunnerAddr = liveRunnerAddr + lp.orchestrator.(*stubOrchestrator).serviceURI = "https://public.example.com" + + challenge, oInfo := requestScopePaymentChallenge(t, lp) + + require.Equal(t, "https://public.example.com", challenge.Orchestrator) + require.Equal(t, challenge.Orchestrator, oInfo.GetTranscoder()) +} + +func TestScopeUsesPublicURLsInResponseWhenLiveRunnerAddrIsConfigured(t *testing.T) { + lp := newScopeOffchainHTTP(t) + liveRunnerAddr, err := url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + lp.node.LiveRunnerAddr = liveRunnerAddr + lp.orchestrator.(*stubOrchestrator).serviceURI = "https://public.example.com" + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp worker.LiveVideoToVideoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotNil(t, resp.ControlUrl) + require.True(t, strings.HasPrefix(*resp.ControlUrl, "https://public.example.com/")) + require.NotNil(t, resp.EventsUrl) + require.True(t, strings.HasPrefix(*resp.EventsUrl, "https://public.example.com/")) + require.NotNil(t, resp.ManifestId) + closeScopeEvents(t, lp, *resp.ManifestId) +} + +func TestScopePaidRetryUsesChallengeManifestID(t *testing.T) { + lp := newScopeHTTP(t) + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(100, 1) + challenge, oInfo := requestScopePaymentChallenge(t, lp) + headers := liveRunnerReservationPaymentHeaders(t, orch, oInfo.GetAuthToken(), challenge.ManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp worker.LiveVideoToVideoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotNil(t, resp.ManifestId) + require.Equal(t, challenge.ManifestID, *resp.ManifestId) + require.NotNil(t, resp.ControlUrl) + require.Contains(t, *resp.ControlUrl, *resp.ManifestId+"-control") + require.NotNil(t, resp.EventsUrl) + require.Contains(t, *resp.EventsUrl, *resp.ManifestId+"-events") + closeScopeEvents(t, lp, *resp.ManifestId) + + balance := orch.Balance(orch.Address(), core.ManifestID(challenge.ManifestID)) + require.NotNil(t, balance) + require.Equal(t, "100", balance.FloatString(0)) + + orch.balanceMu.Lock() + balanceBuckets := make(map[core.ManifestID]*big.Rat, len(orch.balances[orch.Address()])) + for manifestID, balance := range orch.balances[orch.Address()] { + balanceBuckets[manifestID] = balance + } + orch.balanceMu.Unlock() + require.Len(t, balanceBuckets, 1) + require.Contains(t, balanceBuckets, core.ManifestID(challenge.ManifestID)) +} + +func TestScopePaidRetryRecurringAccountingUsesChallengeManifestID(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + lp := newScopeHTTP(t) + lp.node.LivePaymentInterval = time.Second + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(3, 1) + + challenge, oInfo := requestScopePaymentChallenge(t, lp) + headers := liveRunnerReservationPaymentHeadersWithPrice(t, orch, oInfo.GetAuthToken(), challenge.ManifestID, liveRunnerTestPricePerSecond(1)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp worker.LiveVideoToVideoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotNil(t, resp.ManifestId) + require.Equal(t, challenge.ManifestID, *resp.ManifestId) + + balance := orch.Balance(orch.Address(), core.ManifestID(challenge.ManifestID)) + require.NotNil(t, balance) + require.Equal(t, "3", balance.FloatString(0)) + + time.Sleep(time.Second) + eventsCh := trickle.NewLocalPublisher(lp.trickleSrv, challenge.ManifestID+"-events", "application/json") + require.NoError(t, eventsCh.Write(strings.NewReader(`{"event":"tick"}`))) + synctest.Wait() + + balance = orch.Balance(orch.Address(), core.ManifestID(challenge.ManifestID)) + require.NotNil(t, balance) + require.Equal(t, "2", balance.FloatString(0)) + + orch.balanceMu.Lock() + balanceBuckets := make(map[core.ManifestID]*big.Rat, len(orch.balances[orch.Address()])) + for manifestID, balance := range orch.balances[orch.Address()] { + balanceBuckets[manifestID] = balance + } + orch.balanceMu.Unlock() + require.Len(t, balanceBuckets, 1) + require.Contains(t, balanceBuckets, core.ManifestID(challenge.ManifestID)) + + require.NoError(t, eventsCh.Close()) + synctest.Wait() + }) +} + +func TestScopePaidRetryAllowsSegmentManifestToDifferFromAuthToken(t *testing.T) { + lp := newScopeHTTP(t) + orch := lp.orchestrator.(*stubOrchestrator) + orch.balances = make(map[ethcommon.Address]map[core.ManifestID]*big.Rat) + orch.paymentCredit = big.NewRat(100, 1) + _, oInfo := requestScopePaymentChallenge(t, lp) + legacyManifestID := "different-manifest" + headers := liveRunnerReservationPaymentHeaders(t, orch, oInfo.GetAuthToken(), legacyManifestID) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp worker.LiveVideoToVideoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotNil(t, resp.ManifestId) + require.Equal(t, legacyManifestID, *resp.ManifestId) + require.NotNil(t, resp.ControlUrl) + require.Contains(t, *resp.ControlUrl, legacyManifestID+"-control") + require.NotNil(t, resp.EventsUrl) + require.Contains(t, *resp.EventsUrl, legacyManifestID+"-events") + closeScopeEvents(t, lp, legacyManifestID) + + balance := orch.Balance(orch.Address(), core.ManifestID(legacyManifestID)) + require.NotNil(t, balance) + require.Equal(t, "100", balance.FloatString(0)) + require.Nil(t, orch.Balance(orch.Address(), core.ManifestID(oInfo.GetAuthToken().GetSessionId()))) +} + +func TestScopeServerlessOffchainDoesNotRequirePayment(t *testing.T) { + lp := newScopeOffchainHTTP(t) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp worker.LiveVideoToVideoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotNil(t, resp.ManifestId) + require.NotEmpty(t, *resp.ManifestId) + require.NotNil(t, resp.ControlUrl) + require.Contains(t, *resp.ControlUrl, *resp.ManifestId+"-control") + require.NotNil(t, resp.EventsUrl) + require.Contains(t, *resp.EventsUrl, *resp.ManifestId+"-events") + closeScopeEvents(t, lp, *resp.ManifestId) +} + +func TestScopeRejectsOversizedPayload(t *testing.T) { + lp := newScopeHTTP(t) + + oversizedBody := `{"model_id":"scope","gateway_request_id":"` + strings.Repeat("a", maxScopeRequestBodySize) + `"}` + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(oversizedBody)) + setRequestHeaders(req, liveRunnerSenderHeaders(lp.orchestrator.(*stubOrchestrator))) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + require.Contains(t, w.Body.String(), "http: request body too large") +} + +func TestScopePaymentChallengeRejectsInvalidPayer(t *testing.T) { + lp := newScopeHTTP(t) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + req.Header.Set(liveRunnerSenderHeader, "not-an-address") + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusPaymentRequired, w.Code) + require.Contains(t, w.Body.String(), "invalid live runner payment signer address") +} + +func TestLiveRunnerReserveSessionNoCapacity(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusServiceUnavailable, w.Code) +} + +func TestLiveRunnerStopSessionReleasesCapacity(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/stop", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNoContent, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) +} + +func TestLiveRunnerInternalStopSessionRequiresAuth(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID+"/stop", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID+"/stop", nil) + req.Header.Set("Livepeer-Session-Token", "wrong-secret") + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusServiceUnavailable, w.Code) +} + +func TestLiveRunnerInternalStopSessionReleasesCapacity(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/stop", "") + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNoContent, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) +} + +func TestLiveRunnerInternalStopSessionRejectsTokenForDifferentSession(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{Capacity: 2}) + sessionID1 := reserveLiveRunnerSession(t, lp, "runner-1") + sessionID2 := reserveLiveRunnerSession(t, lp, "runner-1") + require.NotEqual(t, sessionID1, sessionID2) + + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + token1, err := manager.SessionTokenForSession("runner-1", sessionID1) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID2+"/stop", nil) + req.Header.Set("Livepeer-Session-Token", token1) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusServiceUnavailable, w.Code) +} + +func TestLiveRunnerProxyForwardsGetPathQueryAndSessionHeaders(t *testing.T) { + var gotPath, gotQuery, gotRunnerRoute, gotSessionID, gotSessionToken, gotSessionControl string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotQuery = r.URL.RawQuery + gotRunnerRoute = r.Header.Get("Livepeer-Runner-Route") + gotSessionID = r.Header.Get("Livepeer-Session-Id") + gotSessionToken = r.Header.Get("Livepeer-Session-Token") + gotSessionControl = r.Header.Get("Livepeer-Session-Control") + w.Header().Set("X-Upstream", "ok") + w.WriteHeader(http.StatusAccepted) + _, _ = w.Write([]byte("proxied")) + })) + defer upstream.Close() + + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{RunnerURL: upstream.URL + "/base"}) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/session/"+sessionID+"/app/v1/foo/bar?x=1&y=two", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusAccepted, w.Code) + require.Equal(t, "proxied", w.Body.String()) + require.Equal(t, "ok", w.Header().Get("X-Upstream")) + require.Equal(t, "/base/v1/foo/bar", gotPath) + require.Equal(t, "x=1&y=two", gotQuery) + require.Equal(t, "runner-1", gotRunnerRoute) + require.Equal(t, sessionID, gotSessionID) + require.NotEmpty(t, gotSessionToken) + require.Equal(t, "http://localhost:1234/runner/runner-1/session/"+sessionID, gotSessionControl) +} + +func TestLiveRunnerProxySessionControlBaseURL(t *testing.T) { + for _, tc := range []struct { + name string + liveRunnerAddr string + wantBase string + }{ + {"falls back to serviceURI", "", "http://localhost:1234"}, + {"uses liveRunnerAddr", "http://go-livepeer:8935", "http://go-livepeer:8935"}, + } { + t.Run(tc.name, func(t *testing.T) { + var gotSessionControl string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSessionControl = r.Header.Get("Livepeer-Session-Control") + w.WriteHeader(http.StatusAccepted) + })) + defer upstream.Close() + + lp := newLiveRunnerHTTP(t, true) + if tc.liveRunnerAddr != "" { + addr, err := url.Parse(tc.liveRunnerAddr) + require.NoError(t, err) + lp.node.LiveRunnerAddr = addr + } + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{RunnerURL: upstream.URL}) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/session/"+sessionID+"/app/x", nil) + lp.ServeHTTP(httptest.NewRecorder(), req) + + require.Equal(t, tc.wantBase+"/runner/runner-1/session/"+sessionID, gotSessionControl) + }) + } +} + +func TestLiveRunnerProxyForwardsPostBody(t *testing.T) { + var gotMethod, gotBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotMethod = r.Method + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + gotBody = string(body) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{RunnerURL: upstream.URL}) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/app/generate", strings.NewReader(`{"prompt":"hi"}`)) + req.Header.Set("Content-Type", "application/json") + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, http.MethodPost, gotMethod) + require.Equal(t, `{"prompt":"hi"}`, gotBody) +} + +func TestLiveRunnerSingleShotProxyForwardsAndReleasesCapacity(t *testing.T) { + var gotPath, gotQuery, gotRunnerRoute, gotSessionID, gotSessionToken, gotSessionControl, gotMethod, gotBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotQuery = r.URL.RawQuery + gotRunnerRoute = r.Header.Get("Livepeer-Runner-Route") + gotSessionID = r.Header.Get("Livepeer-Session-Id") + gotSessionToken = r.Header.Get("Livepeer-Session-Token") + gotSessionControl = r.Header.Get("Livepeer-Session-Control") + gotMethod = r.Method + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + gotBody = string(body) + w.Header().Set("X-Upstream", "ok") + w.WriteHeader(http.StatusAccepted) + _, _ = w.Write([]byte("single-shot")) + })) + defer upstream.Close() + + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{RunnerURL: upstream.URL + "/base", Mode: runner.LiveRunnerModeSingleShot}) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/runner-1/app/v1/foo/bar?x=1&y=two", strings.NewReader(`{"prompt":"hi"}`)) + req.Header.Set("Content-Type", "application/json") + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusAccepted, w.Code) + require.Equal(t, "single-shot", w.Body.String()) + require.Equal(t, "ok", w.Header().Get("X-Upstream")) + require.Equal(t, "/base/v1/foo/bar", gotPath) + require.Equal(t, "x=1&y=two", gotQuery) + require.Equal(t, "runner-1", gotRunnerRoute) + require.NotEmpty(t, gotSessionID) + require.NotEmpty(t, gotSessionToken) + require.Equal(t, "http://localhost:1234/runner/runner-1/session/"+gotSessionID, gotSessionControl) + require.Equal(t, http.MethodPost, gotMethod) + require.Equal(t, `{"prompt":"hi"}`, gotBody) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/apps/runner-1/app/next", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusAccepted, w.Code) +} + +func TestLiveRunnerSingleShotProxyNoCapacityWhileRequestInFlight(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(started) + <-release + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{RunnerURL: upstream.URL, Mode: runner.LiveRunnerModeSingleShot}) + + firstDone := make(chan int, 1) + go func() { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/app/hold", nil) + lp.ServeHTTP(w, req) + firstDone <- w.Code + }() + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("timed out waiting for first single-shot request") + } + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/app/blocked", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusServiceUnavailable, w.Code) + + close(release) + require.Equal(t, http.StatusOK, <-firstDone) +} + +func TestLiveRunnerSingleShotProxyRejectsPersistentRunner(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/app/v1/foo", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + require.Contains(t, w.Body.String(), "single-shot") +} + +func TestLiveRunnerProxyRejectsInvalidSession(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/apps/runner-1/session/missing/app/v1/foo", nil) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusNotFound, w.Code) +} + +func TestLiveRunnerCreateTrickleChannel(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"foo/bar","mime_type":"video/MP2T"},{"name":"events","mime_type":"application/json"}]}`) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerTrickleChannelsResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp.Channels, 2) + require.Equal(t, "foo-bar", resp.Channels[0].Name) + require.Equal(t, sessionID+"-foo-bar", resp.Channels[0].ChannelName) + require.Equal(t, "video/MP2T", resp.Channels[0].MimeType) + require.Equal(t, "http://localhost:1234/ai/trickle/"+sessionID+"-foo-bar", resp.Channels[0].URL) + require.Equal(t, "http://localhost:1234/ai/trickle/"+sessionID+"-foo-bar", resp.Channels[0].InternalURL) + require.Equal(t, sessionID+"-events", resp.Channels[1].ChannelName) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/ai/trickle/"+sessionID+"-foo-bar/next", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) +} + +func TestLiveRunnerCreateTrickleChannelUsesLiveRunnerAddr(t *testing.T) { + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + node.LiveRunnerAddr, err = url.Parse("http://go-livepeer:8935") + require.NoError(t, err) + orch := newStubOrchestrator() + orch.secret = "live-runner-secret" + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}}) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + require.NoError(t, startAIServer(lp)) + + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"override"}]}`) + lp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerTrickleChannelsResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp.Channels, 1) + require.Equal(t, "http://localhost:1234/ai/trickle/"+sessionID+"-override", resp.Channels[0].URL) + require.Equal(t, "http://go-livepeer:8935/ai/trickle/"+sessionID+"-override", resp.Channels[0].InternalURL) +} + +func TestLiveRunnerCreateTrickleChannelReturnsExisting(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"existing","mime_type":"video/MP2T"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var first liveRunnerTrickleChannelsResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &first)) + + w = httptest.NewRecorder() + req = newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"existing","mime_type":"application/json"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var second liveRunnerTrickleChannelsResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &second)) + require.Equal(t, first, second) + require.NotEmpty(t, second.Channels[0].URL) + require.NotEmpty(t, second.Channels[0].InternalURL) + require.Equal(t, "video/MP2T", second.Channels[0].MimeType) +} + +func TestLiveRunnerCreateTrickleChannelRejectsInvalidSessionAndName(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/missing/channels", `{"channels":[{"name":"valid"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) + + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + w = httptest.NewRecorder() + req = newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"bad.name"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestLiveRunnerTrickleChannelRequiresAuth(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + // Check missing session token. + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", strings.NewReader(`{"channels":[{"name":"unauthorized"}]}`)) + req.Header.Set("Authorization", "wrong-secret") + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + + // Check delete requires session token too. + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodDelete, "/runner/runner-1/session/"+sessionID+"/channels", strings.NewReader(`{"channels":["unauthorized"]}`)) + req.Header.Set("Authorization", "wrong-secret") + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) +} + +func TestLiveRunnerTrickleChannelRejectsTokenForDifferentSession(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, &liveRunnerRegistrationOptions{Capacity: 2}) + sessionID1 := reserveLiveRunnerSession(t, lp, "runner-1") + sessionID2 := reserveLiveRunnerSession(t, lp, "runner-1") + require.NotEqual(t, sessionID1, sessionID2) + + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + token1, err := manager.SessionTokenForSession("runner-1", sessionID1) + require.NoError(t, err) + + // Check token is scoped to its session. + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID2+"/channels", strings.NewReader(`{"channels":[{"name":"wrong-session"}]}`)) + req.Header.Set("Livepeer-Session-Token", token1) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) +} + +func TestLiveRunnerTrickleChannelBatchLimit(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + createReq := liveRunnerTrickleChannelsRequest{Channels: make([]liveRunnerTrickleChannelRequest, maxLiveRunnerTrickleChannelsPerRequest+1)} + deleteReq := deleteLiveRunnerTrickleChannelsRequest{Channels: make([]string, maxLiveRunnerTrickleChannelsPerRequest+1)} + for i := range createReq.Channels { + name := fmt.Sprintf("channel-%d", i) + createReq.Channels[i] = liveRunnerTrickleChannelRequest{Name: name} + deleteReq.Channels[i] = name + } + + body, err := json.Marshal(createReq) + require.NoError(t, err) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", bytes.NewReader(body)) + setLiveRunnerSessionToken(t, lp, req, "runner-1", sessionID) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusBadRequest, w.Code) + + body, err = json.Marshal(deleteReq) + require.NoError(t, err) + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodDelete, "/runner/runner-1/session/"+sessionID+"/channels", bytes.NewReader(body)) + setLiveRunnerSessionToken(t, lp, req, "runner-1", sessionID) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestLiveRunnerDeleteTrickleChannel(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"delete-me"},{"name":"delete-me-too"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = newLiveRunnerChannelRequest(lp, http.MethodDelete, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":["delete-me","delete-me-too"]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var resp deleteLiveRunnerTrickleChannelsResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, []string{"delete-me", "delete-me-too"}, resp.Deleted) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/ai/trickle/"+sessionID+"-delete-me/next", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) +} + +func TestLiveRunnerStopSessionClosesTrickleChannels(t *testing.T) { + lp := newLiveRunnerHTTP(t, true) + registerLiveRunnerForSession(t, lp, nil) + sessionID := reserveLiveRunnerSession(t, lp, "runner-1") + + w := httptest.NewRecorder() + req := newLiveRunnerChannelRequest(lp, http.MethodPost, "/runner/runner-1/session/"+sessionID+"/channels", `{"channels":[{"name":"stop-cleanup"}]}`) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/apps/runner-1/session/"+sessionID+"/stop", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNoContent, w.Code) + + w = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/ai/trickle/"+sessionID+"-stop-cleanup/next", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) +} + +func newLiveRunnerHTTP(t *testing.T, withManager bool) *lphttp { + t.Helper() + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + orch := newStubOrchestrator() + orch.secret = "live-runner-secret" + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + if withManager { + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}}) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + } + require.NoError(t, startAIServer(lp)) + return lp +} + +func newLiveRunnerHTTPOnchain(t *testing.T) *lphttp { + t.Helper() + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + node.LivePaymentInterval = 5 * time.Second + orch := newStubOrchestrator() + orch.secret = "live-runner-secret" + orch.ticketParams = defaultTicketParams() + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}, Onchain: true}) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + require.NoError(t, startAIServer(lp)) + return lp +} + +func newServerlessLiveRunnerHTTP(t *testing.T, withManager bool, capacity int) *lphttp { + t.Helper() + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + orch := newStubOrchestrator() + orch.secret = "live-runner-secret" + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + if withManager { + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: testLiveRunnerHost{orchestrator: lp.orchestrator}}) + t.Cleanup(manager.Stop) + node.LiveRunnerManager = manager + } + serverlessWorker, err := worker.NewServerlessWorker("ws://serverless.example.com/ws", capacity) + require.NoError(t, err) + node.AIWorker = serverlessWorker + node.Capabilities = createStubAIWorkerCapabilitiesForPipelineModelId("live-video-to-video", "scope") + require.NoError(t, startAIServer(lp)) + return lp +} + +func newScopeHTTP(t *testing.T) *lphttp { + t.Helper() + oldStorage := drivers.NodeStorage + drivers.NodeStorage = drivers.NewMemoryDriver(nil) + t.Cleanup(func() { drivers.NodeStorage = oldStorage }) + node, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + node.Eth = ð.StubClient{} + node.Capabilities = createStubAIWorkerCapabilitiesForPipelineModelId("live-video-to-video", "scope") + orch := newStubOrchestrator() + orch.ticketParams = defaultTicketParams() + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + require.NoError(t, startAIServer(lp)) + return lp +} + +func newScopeOffchainHTTP(t *testing.T) *lphttp { + t.Helper() + lp := newScopeHTTP(t) + lp.node.Eth = nil + return lp +} + +func newLiveRunnerHTTPWithNode(t *testing.T, node *core.LivepeerNode) *lphttp { + t.Helper() + orch := newStubOrchestrator() + orch.secret = "live-runner-secret" + lp := &lphttp{ + orchestrator: orch, + node: node, + transRPC: http.NewServeMux(), + } + require.NoError(t, startAIServer(lp)) + return lp +} + +func requestLiveRunnerPaymentChallenge(t *testing.T, lp *lphttp, runnerID string) (liveRunnerPaymentChallengeResponse, *lpnet.OrchestratorInfo) { + t.Helper() + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/"+runnerID+"/session", nil) + setRequestHeaders(req, liveRunnerSenderHeaders(lp.orchestrator.(*stubOrchestrator))) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusPaymentRequired, w.Code) + return decodeLiveRunnerPaymentChallenge(t, w.Body.Bytes()) +} + +func requestScopePaymentChallenge(t *testing.T, lp *lphttp) (liveRunnerPaymentChallengeResponse, *lpnet.OrchestratorInfo) { + t.Helper() + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/scope", strings.NewReader(`{"model_id":"scope"}`)) + setRequestHeaders(req, liveRunnerSenderHeaders(lp.orchestrator.(*stubOrchestrator))) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusPaymentRequired, w.Code) + return decodeLiveRunnerPaymentChallenge(t, w.Body.Bytes()) +} + +func closeScopeEvents(t *testing.T, lp *lphttp, manifestID string) { + t.Helper() + eventsCh := trickle.NewLocalPublisher(lp.trickleSrv, manifestID+"-events", "application/json") + require.NoError(t, eventsCh.Close()) +} + +func decodeLiveRunnerPaymentChallenge(t *testing.T, body []byte) (liveRunnerPaymentChallengeResponse, *lpnet.OrchestratorInfo) { + t.Helper() + var challenge liveRunnerPaymentChallengeResponse + require.NoError(t, json.Unmarshal(body, &challenge)) + require.NotEmpty(t, challenge.PaymentParams) + raw, err := base64.StdEncoding.DecodeString(challenge.PaymentParams) + require.NoError(t, err) + var oInfo lpnet.OrchestratorInfo + require.NoError(t, proto.Unmarshal(raw, &oInfo)) + return challenge, &oInfo +} + +func liveRunnerSenderHeaders(orch *stubOrchestrator) http.Header { + headers := http.Header{} + headers.Set(liveRunnerSenderHeader, orch.Address().Hex()) + return headers +} + +func setRequestHeaders(req *http.Request, headers http.Header) { + for k, values := range headers { + for _, v := range values { + req.Header.Add(k, v) + } + } +} + +func liveRunnerReservationPaymentHeaders(t *testing.T, orch *stubOrchestrator, authToken *lpnet.AuthToken, manifestID string) http.Header { + t.Helper() + return liveRunnerReservationPaymentHeadersWithPrice(t, orch, authToken, manifestID, &lpnet.PriceInfo{PricePerUnit: 10, PixelsPerUnit: 1}) +} + +func liveRunnerReservationPaymentHeadersWithPrice(t *testing.T, orch *stubOrchestrator, authToken *lpnet.AuthToken, manifestID string, priceInfo *lpnet.PriceInfo) http.Header { + t.Helper() + md := &core.SegTranscodingMetadata{ + ManifestID: core.ManifestID(manifestID), + Caps: core.NewCapabilities(nil, nil), + AuthToken: authToken, + } + sig, err := orch.Sign(md.Flatten()) + require.NoError(t, err) + segData, err := core.NetSegData(md) + require.NoError(t, err) + segData.Sig = sig + segBytes, err := proto.Marshal(segData) + require.NoError(t, err) + + payment := &lpnet.Payment{ + Sender: orch.Address().Bytes(), + ExpectedPrice: priceInfo, + } + paymentBytes, err := proto.Marshal(payment) + require.NoError(t, err) + + headers := http.Header{} + headers.Set(segmentHeader, base64.StdEncoding.EncodeToString(segBytes)) + headers.Set(paymentHeader, base64.StdEncoding.EncodeToString(paymentBytes)) + return headers +} + +func liveRunnerTestPricePerSecond(pricePerSecond int64) *lpnet.PriceInfo { + return &lpnet.PriceInfo{ + PricePerUnit: pricePerSecond, + PixelsPerUnit: int64(defaultSegInfo.Height) * int64(defaultSegInfo.Width) * int64(defaultSegInfo.FPS), + } +} + +type liveRunnerRegistrationOptions struct { + RunnerID string + RunnerURL string + Capacity int + Mode string +} + +func registerLiveRunnerForSession(t *testing.T, lp *lphttp, opts *liveRunnerRegistrationOptions) { + t.Helper() + if lp == nil { + require.FailNow(t, "live runner lphttp is required") + } + if opts == nil { + opts = &liveRunnerRegistrationOptions{} + } + runnerID := opts.RunnerID + if runnerID == "" { + runnerID = "runner-1" + } + if runnerID != strings.TrimSpace(runnerID) || strings.Contains(runnerID, "/") { + require.FailNow(t, "invalid live runner ID", "runnerID=%q", runnerID) + } + runnerURL := opts.RunnerURL + if runnerURL == "" { + runnerURL = "https://runner.example.com" + } + if strings.TrimSpace(runnerURL) == "" { + require.FailNow(t, "live runner URL is required") + } + capacity := opts.Capacity + if capacity == 0 { + capacity = 1 + } + if capacity < 0 { + require.FailNow(t, "live runner capacity must be non-negative", "capacity=%d", capacity) + } + if opts.Mode != "" && opts.Mode != runner.LiveRunnerModePersistent && opts.Mode != runner.LiveRunnerModeSingleShot { + require.FailNow(t, "invalid live runner mode", "mode=%q", opts.Mode) + } + manager, ok := lp.liveRunnerManager() + require.True(t, ok) + _, err := manager.Heartbeat(runner.LiveRunnerHeartbeatRequest{ + RunnerID: runnerID, + RunnerURL: runnerURL, + Status: "ready", + Mode: opts.Mode, + App: "live-video-to-video/scope", + Capacity: capacity, + PriceInfo: runner.LiveRunnerPriceInfo{ + PricePerUnit: 10, + PixelsPerUnit: 1, + Unit: "WEI", + }, + }, lp.orchestrator.RegistrationSecret()) + require.NoError(t, err) +} + +func reserveLiveRunnerSession(t *testing.T, lp *lphttp, runnerID string) string { + t.Helper() + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/"+runnerID+"/session", nil) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerSessionResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.NotEmpty(t, resp.SessionID) + return resp.SessionID +} + +func reservePaidLiveRunnerSession(t *testing.T, lp *lphttp, runnerID string, priceInfo *lpnet.PriceInfo) string { + t.Helper() + challenge, oInfo := requestLiveRunnerPaymentChallenge(t, lp, runnerID) + headers := liveRunnerReservationPaymentHeadersWithPrice(t, lp.orchestrator.(*stubOrchestrator), oInfo.GetAuthToken(), challenge.ManifestID, priceInfo) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/apps/"+runnerID+"/session", nil) + setRequestHeaders(req, headers) + lp.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + var resp liveRunnerSessionResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, challenge.ManifestID, resp.SessionID) + return resp.SessionID +} + +func newLiveRunnerChannelRequest(lp *lphttp, method, target, body string) *http.Request { + req := httptest.NewRequest(method, target, strings.NewReader(body)) + path := strings.Split(strings.TrimPrefix(target, "/"), "?")[0] + parts := strings.Split(path, "/") + if len(parts) >= 5 && parts[0] == "runner" && parts[2] == "session" { + setLiveRunnerSessionToken(nil, lp, req, parts[1], parts[3]) + } + return req +} + +func setLiveRunnerSessionToken(t require.TestingT, lp *lphttp, req *http.Request, runnerID, sessionID string) { + manager, ok := lp.liveRunnerManager() + if !ok { + if t != nil { + require.FailNow(t, "live runner manager unavailable") + } + return + } + token, err := manager.SessionTokenForSession(runnerID, sessionID) + if err != nil { + return + } + req.Header.Set("Livepeer-Session-Token", token) +} diff --git a/server/ai_session.go b/server/ai_session.go index 8aac600d3b..21bdd15e88 100644 --- a/server/ai_session.go +++ b/server/ai_session.go @@ -511,7 +511,7 @@ func (c *AISessionManager) Select(ctx context.Context, cap core.Capability, mode return sess, err } - clog.V(common.DEBUG).Infof(ctx, "selected orchestrator=%s", sess.Transcoder()) + clog.V(common.DEBUG).Infof(ctx, "selected orchestrator=%s warm=%v", sess.Transcoder(), sess.Warm) return sess, nil } diff --git a/server/cliserver_test.go b/server/cliserver_test.go index 6b3f9800cc..57df172cd0 100644 --- a/server/cliserver_test.go +++ b/server/cliserver_test.go @@ -17,6 +17,7 @@ import ( "testing" "time" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" "github.com/livepeer/go-livepeer/eth" @@ -213,6 +214,65 @@ func TestGetStatus(t *testing.T) { assert.Equal(expected, string(body)) } +func TestRegisterLiveRunnersCLIEndpoint(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + n, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + orch := newStubOrchestrator() + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: orch}) + t.Cleanup(manager.Stop) + n.LiveRunnerManager = manager + s := &LivepeerServer{LivepeerNode: n} + srv := httptest.NewServer(s.cliWebServerHandlers("addr")) + defer srv.Close() + + body := fmt.Sprintf(`{"runners":[{"label":"runner-a","runner_url":"https://runner.example.com","app":"live-video-to-video/scope","capacity":1,"price_info":{"price_per_unit":10,"pixels_per_unit":1,"unit":"WEI"},"health_url":%q}]}`, healthSrv.URL) + res, err := http.Post(srv.URL+"/registerLiveRunners", "application/json", strings.NewReader(body)) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + var reg runner.StaticLiveRunnerRegistrationResponse + require.NoError(t, json.NewDecoder(res.Body).Decode(®)) + require.Len(t, reg.Runners, 1) + require.NotEmpty(t, reg.Runners[0].RunnerID) + require.True(t, reg.Runners[0].Healthy) + + require.Len(t, manager.Runners(), 1) +} + +func TestRegisterLiveRunnersCLIEndpointInvalidBatchDoesNotMutate(t *testing.T) { + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer healthSrv.Close() + + n, err := core.NewLivepeerNode(nil, t.TempDir(), nil) + require.NoError(t, err) + orch := newStubOrchestrator() + manager := runner.NewLiveRunnerRegistry(runner.LiveRunnerRegistryConfig{Host: orch}) + t.Cleanup(manager.Stop) + initialConfig := fmt.Sprintf(`{"runners":[{"label":"runner-a","runner_url":"https://runner.example.com","app":"live-video-to-video/scope","price_info":{"price_per_unit":10,"pixels_per_unit":1,"unit":"WEI"},"health_url":%q}]}`, healthSrv.URL) + _, err = manager.RegisterStaticRunnersJSON([]byte(initialConfig)) + require.NoError(t, err) + n.LiveRunnerManager = manager + s := &LivepeerServer{LivepeerNode: n} + srv := httptest.NewServer(s.cliWebServerHandlers("addr")) + defer srv.Close() + + body := fmt.Sprintf(`{"runners":[{"label":"runner-b","runner_url":"https://runner-two.example.com","app":"live-video-to-video/scope","price_info":{"price_per_unit":10,"pixels_per_unit":1,"unit":"WEI"},"health_url":%q},{"runner_url":"https://runner-three.example.com","app":"live-video-to-video/scope","price_info":{"price_per_unit":10,"pixels_per_unit":1,"unit":"WEI"},"health_url":%q}]}`, healthSrv.URL, healthSrv.URL) + res, err := http.Post(srv.URL+"/registerLiveRunners", "application/json", strings.NewReader(body)) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusBadRequest, res.StatusCode) + require.Len(t, manager.Runners(), 1) + require.Contains(t, manager.Runners()[0].URL, "http://localhost:1234/apps/runner_") + require.Contains(t, manager.Runners()[0].URL, "/session") +} + func TestGetEthChainID(t *testing.T) { require := require.New(t) assert := assert.New(t) diff --git a/server/remote_discovery.go b/server/remote_discovery.go index f336e42c18..bd08efd665 100644 --- a/server/remote_discovery.go +++ b/server/remote_discovery.go @@ -1,16 +1,21 @@ package server import ( + "context" + "encoding/json" "math/big" "net/url" + "reflect" + "slices" "sort" "strings" "sync" "time" + "github.com/livepeer/go-livepeer/ai/runner" + "github.com/livepeer/go-livepeer/clog" "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" - "github.com/livepeer/go-livepeer/net" ) var remoteDiscoveryRefreshInterval = common.WebhookDiscoveryRefreshInterval @@ -46,13 +51,26 @@ type remoteDiscoveryPool struct { } type remoteDiscoveryOrchestrator struct { - OD common.OrchestratorDescriptor + URL *url.URL Capabilities []string + Runners []runner.LiveRunnerDiscoveryRunner +} + +type remoteDiscoveryMergeEntry struct { + orch remoteDiscoveryOrchestrator + runners map[string]runner.LiveRunnerDiscoveryRunner + runnerURLs []string +} + +type remoteDiscoveryEntry struct { + Address string `json:"address"` + Runners []runner.LiveRunnerDiscoveryRunner `json:"runners,omitempty"` } func (o remoteDiscoveryOrchestrator) clone() remoteDiscoveryOrchestrator { orch := o orch.Capabilities = append([]string(nil), o.Capabilities...) + orch.Runners = append([]runner.LiveRunnerDiscoveryRunner(nil), o.Runners...) return orch } @@ -70,7 +88,7 @@ func (p *remoteDiscoveryPool) Orchestrators(caps []string) []remoteDiscoveryOrch seen := make(map[string]bool) for _, key := range caps { for _, orch := range p.caps[key] { - u := orch.OD.LocalInfo.URL.String() + u := orch.URL.String() if seen[u] { continue } @@ -102,67 +120,123 @@ func (p *remoteDiscoveryPool) refresh() { return } - var ods common.OrchestratorDescriptors + var networkCaps []*common.OrchNetworkCapabilities if p.node != nil { - networkCaps := p.node.GetNetworkCapabilities() - ods = make(common.OrchestratorDescriptors, 0, len(networkCaps)) - for _, orchCaps := range networkCaps { - if orchCaps == nil || orchCaps.OrchURI == "" { - continue - } - orchURL, err := url.ParseRequestURI(orchCaps.OrchURI) - if err != nil { - continue - } - ods = append(ods, common.OrchestratorDescriptor{ - LocalInfo: &common.OrchestratorLocalInfo{ - URL: orchURL, - Score: common.Score_Trusted, - }, - RemoteInfo: &net.OrchestratorInfo{ - Transcoder: orchCaps.OrchURI, - Capabilities: orchCaps.Capabilities, - PriceInfo: orchCaps.PriceInfo, - CapabilitiesPrices: orchCaps.CapabilitiesPrices, - Hardware: orchCaps.Hardware, - }, - }) - } - } - - cached := make([]remoteDiscoveryOrchestrator, 0, len(ods)) - caps := make(map[string][]remoteDiscoveryOrchestrator) - for _, od := range ods { - if od.LocalInfo == nil || od.LocalInfo.URL == nil { + networkCaps = p.node.GetNetworkCapabilities() + } + + merged := make(map[string]*remoteDiscoveryMergeEntry) + var order []string + // A node is one discovery `address` value. A single orchestrator may expose + // multiple nodes/addresses through its /discovery response. + validNodes := make(map[string]bool) + rejectedCapabilityNodes := make(map[string]bool) + validApps := make(map[string]map[string]bool) + // Create or reuse the accumulator for a node. The same node address can be + // seen from normal polling and from multiple /discovery responses, so merge + // everything for that address into one final /discover-orchestrators entry. + ensureMerged := func(addr string, orchURL *url.URL) *remoteDiscoveryMergeEntry { + // Treat root service URIs with and without a trailing slash as the same address. + key := strings.TrimRight(addr, "/") + if entry := merged[key]; entry != nil { + return entry + } + entry := &remoteDiscoveryMergeEntry{ + orch: remoteDiscoveryOrchestrator{ + URL: orchURL, + }, + runners: make(map[string]runner.LiveRunnerDiscoveryRunner), + } + merged[key] = entry + order = append(order, key) + return entry + } + + for _, orchCaps := range networkCaps { + if orchCaps == nil || orchCaps.OrchURI == "" { continue } - entry := remoteDiscoveryOrchestrator{ - OD: od, + orchURL, err := url.ParseRequestURI(orchCaps.OrchURI) + if err != nil { + continue } - eligibleCaps := make([]string, 0) + // Treat root service URIs with and without a trailing slash as the same address. + key := strings.TrimRight(orchURL.String(), "/") + eligibleCaps := make(map[string]bool) + advertisedCaps := 0 // Capabilities in discovery are price-eligible capabilities only. - forEachRemoteDiscoveryCapability(od.RemoteInfo, func(key string, capability core.Capability, modelID string) { - price := capabilityPrice(od.RemoteInfo, capability, modelID) + forEachRemoteDiscoveryCapability(orchCaps, func(key string, capability core.Capability, modelID string) { + advertisedCaps++ + price := capabilityPrice(orchCaps, capability, modelID) + if price == nil { + return + } maxPrice := capabilityMaxPrice(capability, modelID) - if maxPrice != nil { - // If a max price is configured, missing price data should not be eligible. - if price == nil || price.Cmp(maxPrice) > 0 { - return - } + if maxPrice != nil && price.Cmp(maxPrice) > 0 { + return } - eligibleCaps = append(eligibleCaps, key) + eligibleCaps[key] = true }) // Drop orchestrators if it doesn't have any eligible capabilities if len(eligibleCaps) == 0 { + if advertisedCaps > 0 { + rejectedCapabilityNodes[key] = true + } continue } - sort.Strings(eligibleCaps) - entry.Capabilities = eligibleCaps - // Keep the capability index aligned with the orchestrator capability list. - for _, key := range eligibleCaps { - caps[key] = append(caps[key], entry.clone()) + validNodes[key] = true + entry := ensureMerged(orchURL.String(), orchURL) + for cap := range eligibleCaps { + entry.orch.Capabilities = append(entry.orch.Capabilities, cap) + } + sort.Strings(entry.orch.Capabilities) + validApps[key] = eligibleCaps + } + + for _, orchCaps := range networkCaps { + if orchCaps == nil { + continue + } + for _, discovery := range remoteDiscoveryEntries(orchCaps.Discovery) { + if discovery.Address == "" { + continue + } + // Treat root service URIs with and without a trailing slash as the same address. + key := strings.TrimRight(discovery.Address, "/") + if validNodes[key] { + mergeDiscoveryRunners(merged[key], discovery, validApps[key]) + continue + } + // If normal polling saw this address but price/capability filtering rejected it, + // do not let runner discovery publish it anyway. + if rejectedCapabilityNodes[key] { + continue + } + orchURL, err := url.ParseRequestURI(discovery.Address) + if err != nil { + continue + } + entry := ensureMerged(discovery.Address, orchURL) + mergeDiscoveryRunners(entry, discovery, nil) + } + } + + cached := make([]remoteDiscoveryOrchestrator, 0, len(order)) + caps := make(map[string][]remoteDiscoveryOrchestrator) + for _, key := range order { + entry := merged[key] + if entry == nil || entry.orch.URL == nil || len(entry.orch.Capabilities) == 0 { + continue + } + sort.Strings(entry.orch.Capabilities) + entry.orch.Runners = make([]runner.LiveRunnerDiscoveryRunner, 0, len(entry.runnerURLs)) + for _, url := range entry.runnerURLs { + entry.orch.Runners = append(entry.orch.Runners, entry.runners[url]) + } + for _, cap := range entry.orch.Capabilities { + caps[cap] = append(caps[cap], entry.orch.clone()) } - cached = append(cached, entry) + cached = append(cached, entry.orch) } // Publish the cache/index snapshot atomically. @@ -179,7 +253,111 @@ func cloneRemoteDiscoveryOrchestrators(cached []remoteDiscoveryOrchestrator) []r return cloned } -func forEachRemoteDiscoveryCapability(info *net.OrchestratorInfo, f func(key string, capability core.Capability, modelID string)) { +func mergeDiscoveryRunners(entry *remoteDiscoveryMergeEntry, discovery remoteDiscoveryEntry, eligibleCaps map[string]bool) { + // For each address, merge runners by URL. First wins; matching duplicates are + // ignored and conflicting duplicates are logged before keeping the first. + for _, r := range discovery.Runners { + if r.URL == "" { + continue + } + if !entry.remoteDiscoveryRunnerEligible(r, eligibleCaps) { + continue + } + if !slices.Contains(entry.orch.Capabilities, r.App) { + entry.orch.Capabilities = append(entry.orch.Capabilities, r.App) + } + if existing, ok := entry.runners[r.URL]; ok { + if !sameRemoteDiscoveryRunner(existing, r) { + clog.Warningf(context.Background(), "Conflicting remote discovery runner for orchestrator=%s runner=%s; keeping first", entry.orch.URL.String(), r.URL) + } + continue + } + entry.runners[r.URL] = r + entry.runnerURLs = append(entry.runnerURLs, r.URL) + } +} + +func sameRemoteDiscoveryRunner(a, b runner.LiveRunnerDiscoveryRunner) bool { + // Capacity usage is expected to move between discovery polls. Ignore it + // when deciding whether duplicate runner URLs represent conflicting metadata. + a.CapacityUsed = 0 + a.CapacityAvailable = 0 + b.CapacityUsed = 0 + b.CapacityAvailable = 0 + return reflect.DeepEqual(a, b) +} + +func remoteDiscoveryEntries(raw json.RawMessage) []remoteDiscoveryEntry { + if len(raw) == 0 { + return nil + } + var entries []remoteDiscoveryEntry + if err := json.Unmarshal(raw, &entries); err != nil { + return nil + } + return entries +} + +func (entry *remoteDiscoveryMergeEntry) remoteDiscoveryRunnerEligible(runner runner.LiveRunnerDiscoveryRunner, eligibleCaps map[string]bool) bool { + if eligibleCaps != nil && !eligibleCaps[runner.App] { + return false + } + if runner.PriceInfo == nil { + return false + } + if runner.PriceInfo.PricePerUnit <= 0 || runner.PriceInfo.PixelsPerUnit <= 0 { + return false + } + if !strings.EqualFold(strings.TrimSpace(runner.PriceInfo.Unit), "WEI") { + clog.Warningf(context.Background(), "Rejecting remote discovery runner with unsupported price unit orch=%s runner=%s app=%s price_per_unit=%d pixels_per_unit=%d unit=%s", + entry.orch.URL.String(), + runner.URL, + runner.App, + runner.PriceInfo.PricePerUnit, + runner.PriceInfo.PixelsPerUnit, + runner.PriceInfo.Unit, + ) + return false + } + + capability, modelID, ok := capabilityModelFromApp(runner.App) + if !ok { + return false + } + maxPrice := capabilityMaxPrice(capability, modelID) + if maxPrice == nil { + return true + } + price := new(big.Rat).SetFrac64(runner.PriceInfo.PricePerUnit, runner.PriceInfo.PixelsPerUnit) + if price.Cmp(maxPrice) > 0 { + clog.Warningf(context.Background(), "Rejecting remote discovery runner above max price orch=%s runner=%s app=%s price_per_unit=%d pixels_per_unit=%d unit=%s price=%s max_price=%s", + entry.orch.URL.String(), + runner.URL, + runner.App, + runner.PriceInfo.PricePerUnit, + runner.PriceInfo.PixelsPerUnit, + runner.PriceInfo.Unit, + price.RatString(), + maxPrice.RatString(), + ) + return false + } + return true +} + +func capabilityModelFromApp(app string) (core.Capability, string, bool) { + pipeline, modelID, ok := strings.Cut(app, "/") + if !ok || pipeline == "" || modelID == "" { + return core.Capability_Unused, "", false + } + capability, err := core.PipelineToCapability(pipeline) + if err != nil { + return core.Capability_Unused, "", false + } + return capability, modelID, true +} + +func forEachRemoteDiscoveryCapability(info *common.OrchNetworkCapabilities, f func(key string, capability core.Capability, modelID string)) { if info == nil || info.Capabilities == nil || info.Capabilities.Constraints == nil { return } @@ -209,21 +387,22 @@ func capabilityToPipeline(capability core.Capability) string { return strings.ToLower(strings.ReplaceAll(name, " ", "-")) } -func capabilityPrice(info *net.OrchestratorInfo, capability core.Capability, modelID string) *big.Rat { - if info != nil { - // Check per-capability price if it exists - for _, capPrice := range info.CapabilitiesPrices { - if capPrice == nil || capPrice.PixelsPerUnit <= 0 || core.Capability(capPrice.Capability) != capability { - continue - } - price := new(big.Rat).SetFrac64(capPrice.PricePerUnit, capPrice.PixelsPerUnit) - if capPrice.Constraint == modelID { - return price - } +func capabilityPrice(info *common.OrchNetworkCapabilities, capability core.Capability, modelID string) *big.Rat { + if info == nil { + return nil + } + // Check per-capability price if it exists + for _, capPrice := range info.CapabilitiesPrices { + if capPrice == nil || capPrice.PixelsPerUnit <= 0 || core.Capability(capPrice.Capability) != capability { + continue + } + price := new(big.Rat).SetFrac64(capPrice.PricePerUnit, capPrice.PixelsPerUnit) + if capPrice.Constraint == modelID { + return price } } // Global fallback if no per-capability price is available. - if info == nil || info.PriceInfo == nil || info.PriceInfo.PixelsPerUnit <= 0 { + if info.PriceInfo == nil || info.PriceInfo.PixelsPerUnit <= 0 { return nil } return new(big.Rat).SetFrac64(info.PriceInfo.PricePerUnit, info.PriceInfo.PixelsPerUnit) diff --git a/server/remote_signer.go b/server/remote_signer.go index f1a3fb548a..8713116bce 100644 --- a/server/remote_signer.go +++ b/server/remote_signer.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math" "math/big" "net/http" "net/url" @@ -17,7 +18,10 @@ import ( ethcommon "github.com/ethereum/go-ethereum/common" "github.com/golang/glog" "github.com/golang/protobuf/proto" + "github.com/livepeer/go-livepeer/ai/runner" + "github.com/livepeer/go-livepeer/byoc" "github.com/livepeer/go-livepeer/clog" + "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" lpcrypto "github.com/livepeer/go-livepeer/crypto" "github.com/livepeer/go-livepeer/monitor" @@ -28,7 +32,9 @@ import ( const HTTPStatusRefreshSession = 480 const HTTPStatusPriceExceeded = 481 const HTTPStatusNoTickets = 482 +const RefreshSessionOrchestratorURLHeader = "Livepeer-Orchestrator-URL" const RemoteType_LiveVideoToVideo = "lv2v" +const RemoteType_BYOC = "byoc" const PipelineLiveVideoToVideo = "live-video-to-video" const remoteSignerAuthIDHeader = "Signer-Auth-Id" @@ -70,11 +76,76 @@ func (ls *LivepeerServer) SignOrchestratorInfo(w http.ResponseWriter, r *http.Re _ = json.NewEncoder(w).Encode(results) } +// SignBYOCJobRequest signs a BYOC job using the V1 binary format (FlattenBYOCJob). +type SignBYOCJobRequestInput struct { + ID string `json:"id"` + Capability string `json:"capability"` + Request string `json:"request"` + Parameters string `json:"parameters"` + TimeoutSeconds int `json:"timeout_seconds"` +} + +type SignBYOCJobRequestResponse struct { + Sender string `json:"sender"` + Signature string `json:"signature"` +} + +func (ls *LivepeerServer) SignBYOCJobRequest(w http.ResponseWriter, r *http.Request) { + ctx := clog.AddVal(r.Context(), "request_id", string(core.RandomManifestID())) + remoteAddr := getRemoteAddr(r) + clog.Info(ctx, "BYOC job signing request", "ip", remoteAddr) + + gw := core.NewBroadcaster(ls.LivepeerNode) + + var req SignBYOCJobRequestInput + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + clog.Errorf(ctx, "Failed to decode SignBYOCJobRequest err=%q", err) + respondJsonError(ctx, w, err, http.StatusBadRequest) + return + } + + if req.ID == "" || req.Capability == "" { + err := fmt.Errorf("sign-byoc-job requires non-empty id and capability") + respondJsonError(ctx, w, err, http.StatusBadRequest) + return + } + if req.TimeoutSeconds <= 0 { + err := fmt.Errorf("sign-byoc-job requires positive timeout_seconds") + respondJsonError(ctx, w, err, http.StatusBadRequest) + return + } + + sigPayload := byoc.FlattenBYOCJob(&byoc.BYOCJobSigningInput{ + ID: req.ID, + Capability: req.Capability, + Request: req.Request, + Parameters: req.Parameters, + TimeoutSeconds: req.TimeoutSeconds, + }) + + sig, err := gw.Sign(sigPayload) + if err != nil { + clog.Errorf(ctx, "Failed to sign BYOC job request err=%q", err) + respondJsonError(ctx, w, err, http.StatusInternalServerError) + return + } + + response := SignBYOCJobRequestResponse{ + Sender: gw.Address().Hex(), + Signature: "0x" + hex.EncodeToString(sig), + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(response) +} + // StartRemoteSignerServer starts the HTTP server for remote signer mode func StartRemoteSignerServer(ls *LivepeerServer, bind string) error { // Register the remote signer endpoints ls.HTTPMux.Handle("POST /sign-orchestrator-info", http.HandlerFunc(ls.SignOrchestratorInfo)) ls.HTTPMux.Handle("POST /generate-live-payment", http.HandlerFunc(ls.GenerateLivePayment)) + ls.HTTPMux.Handle("POST /sign-byoc-job", http.HandlerFunc(ls.SignBYOCJobRequest)) if ls.LivepeerNode.RemoteDiscovery { rdp := RemoteDiscoveryConfig{ Pool: ls.LivepeerNode.OrchestratorPool, @@ -173,11 +244,25 @@ type RemotePaymentRequest struct { // Number of pixels to generate a ticket for. Required if `type` is not set. InPixels int64 `json:"inPixels"` - // Job type to automatically calculate payments. Valid values: `lv2v`. Optional. + // Job type to automatically calculate payments. Valid values: `lv2v`, `byoc`. Optional. Type string `json:"type"` // Capabilities to include in the ticket. Optional; may be set for the lv2v job type. Capabilities []byte `json:"capabilities"` + + // Capability is the real BYOC capability name (e.g. "nano-banana"). When + // set, it overrides the metered `pipeline` label for the emitted + // create_signed_ticket usage event, decoupling usage attribution from + // `Type` (which stays "lv2v" for fee/pixel routing). Optional; empty + // preserves the previous behavior (pipeline falls back to the lv2v + // constant when Type=="lv2v"). + Capability string `json:"capability,omitempty"` + + // ModelID is the specific provider model from the request + // payload/parameters. When set, it overrides the metered `model_id` + // label. Optional; empty preserves the previous behavior (the + // capabilities-derived model id, which defaults to "unknown" downstream). + ModelID string `json:"model_id,omitempty"` } // Returned by the remote signer and includes a new payment plus updated state. @@ -200,7 +285,7 @@ type authResponse struct { // Unix timestamp (seconds) until which auth is considered valid. // Allows for skipping webhook callbacks until this time is exceeded. Expiry int64 `json:"expiry,omitempty"` - // Optional opaque identifier. + // Optional opaque identifier for metering attribution. AuthID string `json:"auth_id,omitempty"` } @@ -286,6 +371,120 @@ func (ls *LivepeerServer) authLivePayment(r *http.Request, state *RemotePaymentS return *webhookResp.Status, &webhookResp, errors.New(webhookResp.Reason) } +// maxUsageLabelLen bounds the length of gateway-supplied usage labels +// (pipeline/model_id) before they are emitted to the metering pipeline. +// req.Capability / req.ModelID come straight from the request body, so a buggy +// or malicious caller could otherwise send very long / high-cardinality values +// that inflate Kafka payload size and downstream label cardinality (the +// pipeline label was previously a small fixed set). 128 runes is generous for +// real capability/model identifiers. +const maxUsageLabelLen = 128 + +// sanitizeUsageLabel trims surrounding whitespace and caps the length (in +// runes, to never emit invalid UTF-8) of a gateway-supplied metering label. +func sanitizeUsageLabel(s string) string { + s = strings.TrimSpace(s) + if r := []rune(s); len(r) > maxUsageLabelLen { + return string(r[:maxUsageLabelLen]) + } + return s +} + +// resolveUsageLabels derives the metering `pipeline` and `model_id` labels for +// the emitted create_signed_ticket usage event. +// +// The real BYOC capability/model supplied by the gateway (req.Capability / +// req.ModelID) take precedence and override the labels whenever they are set, +// regardless of `Type` — this decouples usage attribution from `Type`, which +// still drives fee/pixel routing. The gateway-supplied values are sanitized +// (trimmed + length-capped) to bound payload size and label cardinality. +// +// When those additive fields are empty, the labels fall back to the previous +// behavior: for lv2v the pipeline is the live-video-to-video constant and the +// model id is derived from the request capabilities; for any other request the +// untouched label(s) remain empty (the collector then defaults them to +// "unknown"). This makes non-BYOC (lv2v) callers with no override +// byte-identical to before. +func resolveUsageLabels(req *RemotePaymentRequest, caps *core.Capabilities) (pipeline, modelID string) { + if req.Type == RemoteType_LiveVideoToVideo { + pipeline = PipelineLiveVideoToVideo + if caps != nil { + modelID = caps.ModelIDForCapability(core.Capability_LiveVideoToVideo) + } + } + if req.Type == RemoteType_BYOC && caps != nil { + if m := caps.ModelIDForCapability(core.Capability_BYOC); m != "" { + pipeline = core.CapabilityToPipeline(core.Capability_BYOC) + modelID = m + } + } + if c := sanitizeUsageLabel(req.Capability); c != "" { + pipeline = c + } + if m := sanitizeUsageLabel(req.ModelID); m != "" { + modelID = m + } + return pipeline, modelID +} + +// byocCapabilityName returns the BYOC capability constraint used for per-cap +// pricing lookup. The gateway may supply it as req.Capability (lv2v path) or +// encode it in the BYOC capabilities protobuf (type:"byoc" path). +func byocCapabilityName(req *RemotePaymentRequest, caps *core.Capabilities) string { + if c := sanitizeUsageLabel(req.Capability); c != "" { + return c + } + if caps != nil { + return caps.ModelIDForCapability(core.Capability_BYOC) + } + return "" +} + +func isByocBillingType(t string) bool { + return t == RemoteType_BYOC || t == RemoteType_LiveVideoToVideo +} + +// resolveByocPrice resolves the per-capability BYOC price advertised in the +// orchestrator's OrchestratorInfo.CapabilitiesPrices, keyed on the request +// capability name (req.Capability). The orchestrator appends external BYOC +// capability prices to CapabilitiesPrices as +// {Capability: Capability_BYOC, Constraint: } (see +// core/orchestrator.go GetCapabilitiesPrices), so a match is a BYOC entry whose +// Constraint equals req.Capability. +// +// Granularity is per-capability only: req.ModelID is intentionally NOT used for +// price selection (model_id still flows through for metering attribution). +// +// Returns nil when there is no usable per-capability price (no capability set, +// or no matching entry with a positive rate). A matching entry with a +// non-positive rate is skipped rather than aborting the scan, so a later valid +// duplicate entry for the same constraint (e.g. due to misconfiguration) is +// still honored. Callers fall back to the base oInfo.PriceInfo when nil is +// returned, preserving today's behavior for unconfigured or misconfigured +// capabilities. This is a pure function (no flag, no IO) so it is hermetically +// testable without the CGO/ffmpeg toolchain. +func resolveByocPrice(req *RemotePaymentRequest, oInfo *net.OrchestratorInfo) *net.PriceInfo { + if req == nil || oInfo == nil || req.Capability == "" { + return nil + } + for _, p := range oInfo.CapabilitiesPrices { + if p == nil { + continue + } + if p.Capability != uint32(core.Capability_BYOC) || p.Constraint != req.Capability { + continue + } + // Only honor a valid positive rate. Skip invalid/zero entries and keep + // scanning so a later valid duplicate for the same constraint isn't + // shadowed; if none is found we fall back to base (never zeroing the fee). + if p.PricePerUnit <= 0 || p.PixelsPerUnit <= 0 { + continue + } + return &net.PriceInfo{PricePerUnit: p.PricePerUnit, PixelsPerUnit: p.PixelsPerUnit} + } + return nil +} + // GenerateLivePayment handles remote generation of a payment for live streams. func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Request) { requestID := string(core.RandomManifestID()) @@ -320,7 +519,43 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req respondJsonError(ctx, w, err, http.StatusBadRequest) return } + // BYOC per-capability pricing (flag-gated, default OFF). When enabled and the + // gateway supplied a real capability, resolve the per-capability price from + // the orchestrator's advertised CapabilitiesPrices and use it instead of the + // flat base price. The resolved price is written back to oInfo.PriceInfo so it + // is the single source for: state init, initialPrice, the max-price ceiling, + // the payment's ExpectedPrice (which the orchestrator uses to set its fixed + // per-session price), and the price-doubling guard in validatePrice (which + // reads sess.OrchestratorInfo.PriceInfo) — keeping all of them cap-vs-cap + // consistent. When the flag is OFF or no usable cap price matches, behavior is + // byte-identical to the base-price path. priceInfo := oInfo.PriceInfo + useByocPricing := false + var parsedCaps *core.Capabilities + if len(req.Capabilities) > 0 { + var caps net.Capabilities + if err := proto.Unmarshal(req.Capabilities, &caps); err != nil { + clog.Errorf(ctx, "Failed to unmarshal capabilities err=%q", err) + respondJsonError(ctx, w, err, http.StatusBadRequest) + return + } + parsedCaps = core.CapabilitiesFromNetCapabilities(&caps) + } + capName := byocCapabilityName(&req, parsedCaps) + // BYOC per-cap pricing changes the billing basis (compute-seconds instead + // of lv2v pixels). Apply for lv2v or explicit byoc job types when a + // capability constraint is present. + if ls.LivepeerNode.ByocPerCapPricing && capName != "" && isByocBillingType(req.Type) { + priceReq := req + if priceReq.Capability == "" { + priceReq.Capability = capName + } + if capPrice := resolveByocPrice(&priceReq, &oInfo); capPrice != nil { + priceInfo = capPrice + oInfo.PriceInfo = capPrice + useByocPricing = true + } + } if priceInfo == nil || priceInfo.PricePerUnit == 0 || priceInfo.PixelsPerUnit == 0 { err := fmt.Errorf("missing or zero priceInfo") respondJsonError(ctx, w, err, http.StatusBadRequest) @@ -387,14 +622,8 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req // Embedded within genSegCreds, may be used by orch for payment accounting ManifestID: core.ManifestID(manifestID), } - if len(req.Capabilities) > 0 { - var caps net.Capabilities - if err := proto.Unmarshal(req.Capabilities, &caps); err != nil { - clog.Errorf(ctx, "Failed to unmarshal capabilities err=%q", err) - respondJsonError(ctx, w, err, http.StatusBadRequest) - return - } - streamParams.Capabilities = core.CapabilitiesFromNetCapabilities(&caps) + if parsedCaps != nil { + streamParams.Capabilities = parsedCaps } pmParams := pmTicketParams(oInfo.TicketParams) @@ -443,6 +672,7 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req if should, err := shouldRefreshSession(ctx, sess); err == nil && should { err := errors.New("refresh session for remote signer") + w.Header().Set(RefreshSessionOrchestratorURLHeader, oInfo.Transcoder) respondJsonError(ctx, w, err, HTTPStatusRefreshSession) return } else if err != nil { @@ -458,7 +688,17 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req lastUpdate = now } billableSecs := now.Sub(lastUpdate).Seconds() - if req.Type == RemoteType_LiveVideoToVideo { + if useByocPricing || req.Type == RemoteType_BYOC { + // BYOC per-capability tariff is denominated per compute-second (wei/sec), + // so charge on seconds directly rather than synthesizing lv2v pixels. With + // pixels = ceil(billableSecs) and the resolved per-cap price kept as its + // reduced wei/sec rational, calculateFee yields capPriceWeiPerSec * seconds. + if billableSecs <= 0 { + // preload with 60 seconds, mirroring the lv2v first-call behavior + billableSecs = (60 * time.Second).Seconds() + } + pixels = int64(math.Ceil(billableSecs)) + } else if req.Type == RemoteType_LiveVideoToVideo { info := defaultSegInfo if billableSecs <= 0 { // preload with 60 seconds of data for LV2V @@ -611,7 +851,11 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req sessionStatus = "new" } pipeline := "" - if req.Type == RemoteType_LiveVideoToVideo { + modelID := "" + if streamParams.Capabilities != nil { + pipeline, modelID = streamParams.Capabilities.ConstrainedPipelineModelID() + } + if pipeline == "" && req.Type == RemoteType_LiveVideoToVideo { pipeline = PipelineLiveVideoToVideo } // NB: This could could drop events if tha Kafka queue is full! @@ -619,6 +863,7 @@ func (ls *LivepeerServer) GenerateLivePayment(w http.ResponseWriter, r *http.Req "session_id": state.StateID, "session_status": sessionStatus, "pipeline": pipeline, + "model_id": modelID, "request_id": requestID, "orch_address": orchAddr.Hex(), "orch_url": oInfo.Transcoder, @@ -691,10 +936,14 @@ func GetOrchInfoSig(remoteSignerHost *url.URL, headers map[string]string) (*Orch return &signerResp, nil } +// discoveryResponse is intentionally typed. Do NOT add raw json.RawMessage blobs +// here or pass through arbitrary orchestrator /discovery fields; every exposed +// response field must be reviewed and modeled explicitly. type discoveryResponse struct { - Address string `json:"address,omitempty"` - Score float32 `json:"score,omitempty"` - Capabilities []string `json:"capabilities,omitempty"` + Address string `json:"address,omitempty"` + Score float32 `json:"score,omitempty"` + Capabilities []string `json:"capabilities,omitempty"` + Runners []runner.LiveRunnerDiscoveryRunner `json:"runners,omitempty"` } // GetOrchestrators returns the configured orchestrators in webhook-compatible format @@ -724,11 +973,11 @@ func (ls *LivepeerServer) GetOrchestrators(pool *remoteDiscoveryPool, w http.Res infos := pool.Orchestrators(filteredCaps) resp := make([]discoveryResponse, 0, len(infos)) for _, cached := range infos { - od := cached.OD resp = append(resp, discoveryResponse{ - Address: od.LocalInfo.URL.String(), - Score: od.LocalInfo.Score, + Address: cached.URL.String(), + Score: common.Score_Trusted, // Legacy go-livepeer webhook field. Capabilities: append([]string(nil), cached.Capabilities...), + Runners: append([]runner.LiveRunnerDiscoveryRunner(nil), cached.Runners...), }) } diff --git a/server/remote_signer_test.go b/server/remote_signer_test.go index 375eb82c39..99a0b65a33 100644 --- a/server/remote_signer_test.go +++ b/server/remote_signer_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/http/httptest" "net/url" + "strings" "testing" "testing/synctest" "time" @@ -20,6 +21,7 @@ import ( "github.com/ethereum/go-ethereum/crypto" "github.com/ethereum/go-ethereum/signer/core/apitypes" "github.com/golang/protobuf/proto" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/common" "github.com/livepeer/go-livepeer/core" "github.com/livepeer/go-livepeer/eth" @@ -90,8 +92,9 @@ func TestGenerateLivePayment_RequestValidationErrors(t *testing.T) { defer BroadcastCfg.SetCapabilityMaxPrice(capability1, modelID1, nil) // Clean up baseOrchInfo := &net.OrchestratorInfo{ - Address: ethcommon.HexToAddress("0x1").Bytes(), - PriceInfo: &net.PriceInfo{PricePerUnit: 1, PixelsPerUnit: 1}, + Address: ethcommon.HexToAddress("0x1").Bytes(), + Transcoder: "http://orch.example", + PriceInfo: &net.PriceInfo{PricePerUnit: 1, PixelsPerUnit: 1}, TicketParams: &net.TicketParams{ Recipient: pm.RandAddress().Bytes(), }, @@ -118,6 +121,7 @@ func TestGenerateLivePayment_RequestValidationErrors(t *testing.T) { sender *pm.MockSender wantStatus int wantMsg string + wantHeader string }{ { name: "invalid JSON", @@ -203,6 +207,7 @@ func TestGenerateLivePayment_RequestValidationErrors(t *testing.T) { }(), wantStatus: HTTPStatusRefreshSession, wantMsg: "refresh session for remote signer", + wantHeader: "http://orch.example", }, { name: "ticket params expired triggers 480", @@ -214,6 +219,7 @@ func TestGenerateLivePayment_RequestValidationErrors(t *testing.T) { sender: newMockSender(mockSenderConfig{validateErr: pm.ErrTicketParamsExpired}), wantStatus: HTTPStatusRefreshSession, wantMsg: "refresh session for remote signer", + wantHeader: "http://orch.example", }, { name: "invalid job type", @@ -319,6 +325,9 @@ func TestGenerateLivePayment_RequestValidationErrors(t *testing.T) { require.NoError(json.NewDecoder(rr.Body).Decode(&apiErr)) require.NotEmpty(apiErr.Error.Message) require.Contains(apiErr.Error.Message, tt.wantMsg) + if tt.wantHeader != "" { + require.Equal(tt.wantHeader, rr.Header().Get(RefreshSessionOrchestratorURLHeader)) + } }) } } @@ -1092,6 +1101,39 @@ func TestRemoteSigner_Discovery(t *testing.T) { Constraint: "model-a", }, }, + Discovery: discoveryRaw(t, `[ + { + "address": "https://orch1.example.com:8935/", + "score": 99, + "capabilities": ["evil/app"], + "unknown_field": "must-not-leak", + "runners": [ + {"url":"https://orch1.example.com:8935/scope-ok","app":"live-video-to-video/model-a","price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://orch1.example.com:8935/scope-missing-price","app":"live-video-to-video/model-a"}, + {"url":"https://orch1.example.com:8935/scope-too-expensive","app":"live-video-to-video/model-a","price_info":{"price_per_unit":120,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://orch1.example.com:8935/scope-usd","app":"live-video-to-video/model-a","price_info":{"price_per_unit":1,"pixels_per_unit":1,"unit":"USD"}}, + {"url":"https://orch1.example.com:8935/scope-invalid-price","app":"live-video-to-video/model-a","price_info":{"price_per_unit":1,"pixels_per_unit":0,"unit":"WEI"}}, + {"url":"https://orch1.example.com:8935/scope-missing-unit","app":"live-video-to-video/model-a","price_info":{"price_per_unit":1,"pixels_per_unit":1}}, + {"url":"https://orch1.example.com:8935/txt","app":"text-to-image/model-b","price_info":{"price_per_unit":1,"pixels_per_unit":1,"unit":"WEI"}} + ] + }, + { + "address": "https://discovered.example.com:8935", + "runners": [ + {"url":"https://discovered.example.com:8935/scope","app":"live-video-to-video/model-a","capacity":3,"price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://discovered.example.com:8935/dupe","app":"live-video-to-video/model-a","capacity":1,"price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://discovered.example.com:8935/missing-price","app":"live-video-to-video/model-a"}, + {"url":"https://discovered.example.com:8935/too-expensive","app":"live-video-to-video/model-a","price_info":{"price_per_unit":120,"pixels_per_unit":1,"unit":"WEI"}} + ] + }, + { + "address": "https://discovered.example.com:8935/", + "runners": [ + {"url":"https://discovered.example.com:8935/dupe","app":"live-video-to-video/model-a","capacity":1,"capacity_used":1,"capacity_available":0,"price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://discovered.example.com:8935/dupe","app":"live-video-to-video/model-a","capacity":2,"price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}} + ] + } + ]`), }, { OrchURI: "https://orch2.example.com:8935", @@ -1114,6 +1156,21 @@ func TestRemoteSigner_Discovery(t *testing.T) { Constraint: "model-b", }, }, + Discovery: discoveryRaw(t, `[ + { + "address": "https://orch2.example.com:8935", + "runners": [ + {"url":"https://orch2.example.com:8935/txt","app":"text-to-image/model-b","price_info":{"price_per_unit":999,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://orch2.example.com:8935/scope","app":"live-video-to-video/model-a","price_info":{"price_per_unit":1,"pixels_per_unit":1,"unit":"WEI"}} + ] + }, + { + "address": "https://discovered.example.com:8935", + "runners": [ + {"url":"https://discovered.example.com:8935/from-orch2","app":"live-video-to-video/model-a","price_info":{"price_per_unit":80,"pixels_per_unit":1,"unit":"WEI"}} + ] + } + ]`), }, { OrchURI: "https://orch3.example.com:8935", @@ -1129,6 +1186,12 @@ func TestRemoteSigner_Discovery(t *testing.T) { Constraint: "model-a", }, }, + Discovery: discoveryRaw(t, `[{ + "address": "https://orch3.example.com:8935", + "runners": [ + {"url":"https://orch3.example.com:8935/scope","app":"live-video-to-video/model-a","price_info":{"price_per_unit":1,"pixels_per_unit":1,"unit":"WEI"}} + ] + }]`), }, // Invalid URI should be dropped during refresh and never returned from discovery. { @@ -1153,15 +1216,25 @@ func TestRemoteSigner_Discovery(t *testing.T) { require.Equal(http.StatusOK, rr.Code) require.Equal("application/json", rr.Header().Get("Content-Type")) + body := rr.Body.Bytes() var resp []discoveryResponse - require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) - require.Len(resp, 2) - require.Equal("https://orch1.example.com:8935", resp[0].Address) - require.Greater(resp[0].Score, float32(0)) - require.Equal([]string{"live-video-to-video/model-a"}, resp[0].Capabilities) - require.Equal("https://orch2.example.com:8935", resp[1].Address) - require.Greater(resp[1].Score, float32(0)) - require.Equal([]string{"text-to-image/model-b"}, resp[1].Capabilities) + require.NoError(json.Unmarshal(body, &resp)) + require.Len(resp, 3) + orch1Resp := discoveryResponseByAddress(t, resp, "https://orch1.example.com:8935") + require.Equal(float32(common.Score_Trusted), orch1Resp.Score) + require.Equal([]string{"live-video-to-video/model-a"}, orch1Resp.Capabilities) + require.Equal([]string{"live-video-to-video/model-a"}, discoveryRunnerApps(t, orch1Resp)) + requireNotContainsDiscoveryField(t, body, "unknown_field") + discoveredResp := discoveryResponseByAddress(t, resp, "https://discovered.example.com:8935") + require.Equal(float32(common.Score_Trusted), discoveredResp.Score) + require.Equal([]string{"live-video-to-video/model-a"}, discoveredResp.Capabilities) + require.Equal([]string{"live-video-to-video/model-a", "live-video-to-video/model-a", "live-video-to-video/model-a"}, discoveryRunnerApps(t, discoveredResp)) + require.Equal(1, discoveredResp.Runners[1].Capacity) + require.Equal(0, discoveredResp.Runners[1].CapacityUsed) + orch2Resp := discoveryResponseByAddress(t, resp, "https://orch2.example.com:8935") + require.Equal(float32(common.Score_Trusted), orch2Resp.Score) + require.Equal([]string{"text-to-image/model-b"}, orch2Resp.Capabilities) + require.Equal([]string{"text-to-image/model-b"}, discoveryRunnerApps(t, orch2Resp)) capsReq := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/model-a", nil) capsRR := httptest.NewRecorder() @@ -1169,8 +1242,9 @@ func TestRemoteSigner_Discovery(t *testing.T) { require.Equal(http.StatusOK, capsRR.Code) var capsResp []discoveryResponse require.NoError(json.NewDecoder(capsRR.Body).Decode(&capsResp)) - require.Len(capsResp, 1) - require.Equal("https://orch1.example.com:8935", capsResp[0].Address) + require.Len(capsResp, 2) + require.Equal([]string{"live-video-to-video/model-a"}, discoveryRunnerApps(t, discoveryResponseByAddress(t, capsResp, "https://orch1.example.com:8935"))) + require.Equal([]string{"live-video-to-video/model-a", "live-video-to-video/model-a", "live-video-to-video/model-a"}, discoveryRunnerApps(t, discoveryResponseByAddress(t, capsResp, "https://discovered.example.com:8935"))) repeatedReq := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/model-a&caps=text-to-image/model-b", nil) repeatedRR := httptest.NewRecorder() @@ -1178,9 +1252,10 @@ func TestRemoteSigner_Discovery(t *testing.T) { require.Equal(http.StatusOK, repeatedRR.Code) var repeatedResp []discoveryResponse require.NoError(json.NewDecoder(repeatedRR.Body).Decode(&repeatedResp)) - require.Len(repeatedResp, 2) - require.Equal("https://orch1.example.com:8935", repeatedResp[0].Address) - require.Equal("https://orch2.example.com:8935", repeatedResp[1].Address) + require.Len(repeatedResp, 3) + discoveryResponseByAddress(t, repeatedResp, "https://orch1.example.com:8935") + discoveryResponseByAddress(t, repeatedResp, "https://discovered.example.com:8935") + discoveryResponseByAddress(t, repeatedResp, "https://orch2.example.com:8935") // If refresh only receives invalid or ineligible network capability entries, // cache should remain empty and discovery should return service unavailable. @@ -1237,6 +1312,280 @@ func TestRemoteSigner_Discovery(t *testing.T) { require.Equal("application/json", ineligibleRR.Header().Get("Content-Type")) } +func TestRemoteSigner_Discovery_UsesRunnerDiscoveryWhenRPCCapabilitiesMissing(t *testing.T) { + require := require.New(t) + + capability := core.Capability_LiveVideoToVideo + modelID := "scope" + BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, core.NewFixedPrice(big.NewRat(200, 995328000000))) + defer BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, nil) + + node := &core.LivepeerNode{} + require.NoError(node.UpdateNetworkCapabilities([]*common.OrchNetworkCapabilities{ + { + OrchURI: "https://runner-derived.example.com:8935", + Discovery: discoveryRaw(t, `[{ + "address": "https://runner-derived.example.com:8935", + "runners": [ + {"url":"https://runner-derived.example.com:8935/apps/runner/session","app":"live-video-to-video/scope","capacity":1,"price_info":{"price_per_unit":5,"pixels_per_unit":995328000000,"unit":"WEI"}} + ] + }]`), + }, + })) + + rdp := &remoteDiscoveryPool{ + node: node, + refreshEvery: time.Hour, + } + ls := &LivepeerServer{} + + req := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/scope", nil) + rr := httptest.NewRecorder() + ls.GetOrchestrators(rdp, rr, req) + + require.Equal(http.StatusOK, rr.Code) + var resp []discoveryResponse + require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) + require.Len(resp, 1) + require.Equal("https://runner-derived.example.com:8935", resp[0].Address) + require.Equal([]string{"live-video-to-video/scope"}, resp[0].Capabilities) + require.Equal([]string{"live-video-to-video/scope"}, discoveryRunnerApps(t, resp[0])) +} + +func TestRemoteSigner_Discovery_RunnerDiscoveryKeepsEligibleRunnersWhenRPCCapabilitiesMissing(t *testing.T) { + require := require.New(t) + + capability := core.Capability_LiveVideoToVideo + modelID := "scope" + BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, core.NewFixedPrice(big.NewRat(200, 995328000000))) + defer BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, nil) + + node := &core.LivepeerNode{} + require.NoError(node.UpdateNetworkCapabilities([]*common.OrchNetworkCapabilities{ + { + OrchURI: "https://mixed-runner-derived.example.com:8935", + Discovery: discoveryRaw(t, `[{ + "address": "https://mixed-runner-derived.example.com:8935", + "runners": [ + {"url":"https://mixed-runner-derived.example.com:8935/too-expensive","app":"live-video-to-video/scope","price_info":{"price_per_unit":201,"pixels_per_unit":995328000000,"unit":"WEI"}}, + {"url":"https://mixed-runner-derived.example.com:8935/usd","app":"live-video-to-video/scope","price_info":{"price_per_unit":5,"pixels_per_unit":995328000000,"unit":"USD"}}, + {"url":"https://mixed-runner-derived.example.com:8935/missing-price","app":"live-video-to-video/scope"}, + {"url":"https://mixed-runner-derived.example.com:8935/invalid-pixels","app":"live-video-to-video/scope","price_info":{"price_per_unit":5,"pixels_per_unit":0,"unit":"WEI"}}, + {"url":"https://mixed-runner-derived.example.com:8935/apps/runner/session","app":"live-video-to-video/scope","capacity":1,"price_info":{"price_per_unit":5,"pixels_per_unit":995328000000,"unit":"WEI"}} + ] + }]`), + }, + })) + + rdp := &remoteDiscoveryPool{ + node: node, + refreshEvery: time.Hour, + } + ls := &LivepeerServer{} + + req := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/scope", nil) + rr := httptest.NewRecorder() + ls.GetOrchestrators(rdp, rr, req) + + require.Equal(http.StatusOK, rr.Code) + var resp []discoveryResponse + require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) + require.Len(resp, 1) + require.Equal("https://mixed-runner-derived.example.com:8935", resp[0].Address) + require.Equal([]string{"live-video-to-video/scope"}, resp[0].Capabilities) + require.Len(resp[0].Runners, 1) + require.Equal("https://mixed-runner-derived.example.com:8935/apps/runner/session", resp[0].Runners[0].URL) + require.Equal([]string{"live-video-to-video/scope"}, discoveryRunnerApps(t, resp[0])) +} + +func TestRemoteSigner_Discovery_FiltersRunnerDiscoveryPricing(t *testing.T) { + require := require.New(t) + + capability := core.Capability_LiveVideoToVideo + modelID := "scope" + BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, core.NewFixedPrice(big.NewRat(200, 995328000000))) + defer BroadcastCfg.SetCapabilityMaxPrice(capability, modelID, nil) + + node := &core.LivepeerNode{} + require.NoError(node.UpdateNetworkCapabilities([]*common.OrchNetworkCapabilities{ + { + OrchURI: "https://filtered.example.com:8935", + Discovery: discoveryRaw(t, `[{ + "address": "https://filtered.example.com:8935", + "runners": [ + {"url":"https://filtered.example.com:8935/too-expensive","app":"live-video-to-video/scope","price_info":{"price_per_unit":201,"pixels_per_unit":995328000000,"unit":"WEI"}}, + {"url":"https://filtered.example.com:8935/usd","app":"live-video-to-video/scope","price_info":{"price_per_unit":5,"pixels_per_unit":995328000000,"unit":"USD"}}, + {"url":"https://filtered.example.com:8935/missing-price","app":"live-video-to-video/scope"}, + {"url":"https://filtered.example.com:8935/non-positive-price","app":"live-video-to-video/scope","price_info":{"price_per_unit":0,"pixels_per_unit":995328000000,"unit":"WEI"}}, + {"url":"https://filtered.example.com:8935/non-positive-pixels","app":"live-video-to-video/scope","price_info":{"price_per_unit":5,"pixels_per_unit":0,"unit":"WEI"}} + ] + }]`), + }, + })) + + rdp := &remoteDiscoveryPool{ + node: node, + refreshEvery: time.Hour, + } + ls := &LivepeerServer{} + + req := httptest.NewRequest(http.MethodGet, "/discover-orchestrators", nil) + rr := httptest.NewRecorder() + ls.GetOrchestrators(rdp, rr, req) + + require.Equal(http.StatusServiceUnavailable, rr.Code) + require.Equal("application/json", rr.Header().Get("Content-Type")) +} + +func discoveryRaw(t *testing.T, data string) json.RawMessage { + t.Helper() + require.True(t, json.Valid([]byte(data))) + return json.RawMessage(data) +} + +func discoveryRunnerApps(t *testing.T, resp discoveryResponse) []string { + t.Helper() + apps := make([]string, 0, len(resp.Runners)) + for _, runner := range resp.Runners { + apps = append(apps, runner.App) + } + return apps +} + +func discoveryResponseByAddress(t *testing.T, resp []discoveryResponse, address string) discoveryResponse { + t.Helper() + for _, entry := range resp { + if entry.Address == address { + return entry + } + } + require.Failf(t, "missing discovery response", "address=%s", address) + return discoveryResponse{} +} + +func requireNotContainsDiscoveryField(t *testing.T, body []byte, field string) { + t.Helper() + var raw []map[string]json.RawMessage + require.NoError(t, json.Unmarshal(body, &raw)) + for _, entry := range raw { + require.NotContains(t, entry, field) + } +} + +func TestRemoteDiscoveryRunnerDuplicateComparison(t *testing.T) { + base := runner.LiveRunnerDiscoveryRunner{ + URL: "https://runner.example.com/scope", + App: "live-video-to-video/model-a", + Capacity: 2, + CapacityUsed: 1, + CapacityAvailable: 1, + PriceInfo: &runner.LiveRunnerPriceInfo{ + PricePerUnit: 7, + PixelsPerUnit: 1, + Unit: "WEI", + }, + } + + usageChanged := base + usageChanged.CapacityUsed = 0 + usageChanged.CapacityAvailable = 2 + require.True(t, sameRemoteDiscoveryRunner(base, usageChanged)) + + capacityChanged := base + capacityChanged.Capacity = 3 + require.False(t, sameRemoteDiscoveryRunner(base, capacityChanged)) +} + +func TestRemoteSigner_Discovery_RequiresAdvertisedPriceWithoutMaxPrice(t *testing.T) { + require := require.New(t) + + capability := core.Capability_LiveVideoToVideo + modelPriced := "model-priced" + modelUnpriced := "model-unpriced" + BroadcastCfg.SetMaxPrice(nil) + BroadcastCfg.SetCapabilityMaxPrice(capability, modelPriced, nil) + BroadcastCfg.SetCapabilityMaxPrice(capability, modelUnpriced, nil) + + caps := core.NewCapabilities([]core.Capability{capability}, nil) + caps.SetPerCapabilityConstraints(core.PerCapabilityConstraints{ + capability: &core.CapabilityConstraints{ + Models: map[string]*core.ModelConstraint{ + modelPriced: {}, + modelUnpriced: {}, + }, + }, + }) + + node := &core.LivepeerNode{} + require.NoError(node.UpdateNetworkCapabilities([]*common.OrchNetworkCapabilities{ + { + OrchURI: "https://priced.example.com:8935", + Capabilities: caps.ToNetCapabilities(), + CapabilitiesPrices: []*net.PriceInfo{ + { + PricePerUnit: 7, + PixelsPerUnit: 1, + Capability: uint32(capability), + Constraint: modelPriced, + }, + }, + Discovery: discoveryRaw(t, `[{ + "address": "https://priced.example.com:8935", + "runners": [ + {"url":"https://priced.example.com:8935/priced","app":"live-video-to-video/model-priced","price_info":{"price_per_unit":7,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://priced.example.com:8935/unpriced","app":"live-video-to-video/model-unpriced","price_info":{"price_per_unit":7,"pixels_per_unit":1,"unit":"WEI"}} + ] + }]`), + }, + { + OrchURI: "https://unpriced.example.com:8935", + Capabilities: caps.ToNetCapabilities(), + Discovery: discoveryRaw(t, `[{ + "address": "https://unpriced.example.com:8935", + "runners": [ + {"url":"https://unpriced.example.com:8935/priced","app":"live-video-to-video/model-priced","price_info":{"price_per_unit":7,"pixels_per_unit":1,"unit":"WEI"}}, + {"url":"https://unpriced.example.com:8935/unpriced","app":"live-video-to-video/model-unpriced","price_info":{"price_per_unit":7,"pixels_per_unit":1,"unit":"WEI"}} + ] + }]`), + }, + })) + + rdp := &remoteDiscoveryPool{ + node: node, + refreshEvery: time.Hour, + } + ls := &LivepeerServer{} + + req := httptest.NewRequest(http.MethodGet, "/discover-orchestrators", nil) + rr := httptest.NewRecorder() + ls.GetOrchestrators(rdp, rr, req) + + require.Equal(http.StatusOK, rr.Code) + var resp []discoveryResponse + require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) + require.Len(resp, 1) + require.Equal("https://priced.example.com:8935", resp[0].Address) + require.Equal([]string{"live-video-to-video/model-priced"}, resp[0].Capabilities) + require.Equal([]string{"live-video-to-video/model-priced"}, discoveryRunnerApps(t, resp[0])) + + capsReq := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/model-priced", nil) + capsRR := httptest.NewRecorder() + ls.GetOrchestrators(rdp, capsRR, capsReq) + require.Equal(http.StatusOK, capsRR.Code) + var capsResp []discoveryResponse + require.NoError(json.NewDecoder(capsRR.Body).Decode(&capsResp)) + require.Len(capsResp, 1) + require.Equal("https://priced.example.com:8935", capsResp[0].Address) + + unpricedReq := httptest.NewRequest(http.MethodGet, "/discover-orchestrators?caps=live-video-to-video/model-unpriced", nil) + unpricedRR := httptest.NewRecorder() + ls.GetOrchestrators(rdp, unpricedRR, unpricedReq) + require.Equal(http.StatusOK, unpricedRR.Code) + var unpricedResp []discoveryResponse + require.NoError(json.NewDecoder(unpricedRR.Body).Decode(&unpricedResp)) + require.Empty(unpricedResp) +} + func TestRemoteSigner_Discovery_RefreshesAfterInterval(t *testing.T) { synctest.Test(t, func(t *testing.T) { require := require.New(t) @@ -1344,3 +1693,415 @@ func TestGetOrchInfoSig_SendsConfiguredHeaders(t *testing.T) { require.Equal([]byte{0x12, 0x34}, []byte(resp.Address)) require.Equal([]byte{0xab, 0xcd}, []byte(resp.Signature)) } + +// lv2vCapsWithModel builds a core.Capabilities advertising a single warm model +// for the live-video-to-video capability (used to exercise the existing +// capabilities-derived model_id fallback). +func lv2vCapsWithModel(t *testing.T, modelID string) *core.Capabilities { + t.Helper() + c := core.NewCapabilities([]core.Capability{core.Capability_LiveVideoToVideo}, nil) + c.SetPerCapabilityConstraints(core.PerCapabilityConstraints{ + core.Capability_LiveVideoToVideo: &core.CapabilityConstraints{ + Models: core.ModelConstraints{modelID: &core.ModelConstraint{Warm: true}}, + }, + }) + return c +} + +func byocCapsWithModel(t *testing.T, modelID string) *core.Capabilities { + t.Helper() + c := core.NewCapabilities([]core.Capability{core.Capability_BYOC}, nil) + c.SetPerCapabilityConstraints(core.PerCapabilityConstraints{ + core.Capability_BYOC: &core.CapabilityConstraints{ + Models: core.ModelConstraints{modelID: &core.ModelConstraint{Warm: true}}, + }, + }) + return c +} + +func byocCapsProtoBytes(t *testing.T, modelID string) []byte { + t.Helper() + caps := byocCapsWithModel(t, modelID) + b, err := proto.Marshal(caps.ToNetCapabilities()) + require.NoError(t, err) + return b +} + +// TestResolveUsageLabels asserts the additive usage-attribution contract: +// - when the gateway supplies the real BYOC capability/model, the metered +// pipeline + model_id reflect them (the root-cause fix); +// - when those fields are empty, behavior is unchanged (lv2v constant + +// capabilities-derived model id, or empty → "unknown" downstream). +// +// This is hermetic: it exercises the pure label-resolution helper without the +// CGO ffmpeg toolchain, the remote-signer HTTP handler, or Kafka. +func TestResolveUsageLabels(t *testing.T) { + tests := []struct { + name string + req RemotePaymentRequest + caps *core.Capabilities + wantPipeline string + wantModelID string + }{ + { + name: "lv2v unchanged: no new fields, no caps", + req: RemotePaymentRequest{Type: RemoteType_LiveVideoToVideo}, + caps: nil, + wantPipeline: PipelineLiveVideoToVideo, + wantModelID: "", + }, + { + name: "lv2v unchanged: model derived from capabilities", + req: RemotePaymentRequest{Type: RemoteType_LiveVideoToVideo}, + caps: lv2vCapsWithModel(t, "streamdiffusion"), + wantPipeline: PipelineLiveVideoToVideo, + wantModelID: "streamdiffusion", + }, + { + name: "byoc type: labels from capabilities protobuf", + req: RemotePaymentRequest{ + Type: RemoteType_BYOC, + }, + caps: byocCapsWithModel(t, "flux-schnell"), + wantPipeline: "byoc", + wantModelID: "flux-schnell", + }, + { + name: "byoc: real capability + model override lv2v defaults", + req: RemotePaymentRequest{ + Type: RemoteType_LiveVideoToVideo, + Capability: "nano-banana", + ModelID: "google/nano-banana", + }, + caps: lv2vCapsWithModel(t, "streamdiffusion"), + wantPipeline: "nano-banana", + wantModelID: "google/nano-banana", + }, + { + name: "byoc: capability set, empty model falls back to caps-derived", + req: RemotePaymentRequest{ + Type: RemoteType_LiveVideoToVideo, + Capability: "nano-banana", + }, + caps: lv2vCapsWithModel(t, "streamdiffusion"), + wantPipeline: "nano-banana", + wantModelID: "streamdiffusion", + }, + { + name: "non-lv2v with no fields: both empty (collector defaults to unknown)", + req: RemotePaymentRequest{}, + caps: nil, + wantPipeline: "", + wantModelID: "", + }, + { + name: "byoc: whitespace-only override is ignored, lv2v defaults kept", + req: RemotePaymentRequest{ + Type: RemoteType_LiveVideoToVideo, + Capability: " ", + ModelID: "\t\n", + }, + caps: lv2vCapsWithModel(t, "streamdiffusion"), + wantPipeline: PipelineLiveVideoToVideo, + wantModelID: "streamdiffusion", + }, + { + name: "byoc: oversized override labels are length-capped", + req: RemotePaymentRequest{ + Type: RemoteType_LiveVideoToVideo, + Capability: strings.Repeat("c", maxUsageLabelLen+50), + ModelID: strings.Repeat("m", maxUsageLabelLen+50), + }, + caps: nil, + wantPipeline: strings.Repeat("c", maxUsageLabelLen), + wantModelID: strings.Repeat("m", maxUsageLabelLen), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require := require.New(t) + req := tc.req + pipeline, modelID := resolveUsageLabels(&req, tc.caps) + require.Equal(tc.wantPipeline, pipeline, "pipeline") + require.Equal(tc.wantModelID, modelID, "model_id") + }) + } +} + +// TestResolveByocPrice covers the pure per-capability BYOC price resolver: +// it matches a Capability_BYOC entry whose Constraint equals req.Capability, +// ignores model_id and non-BYOC entries, and returns nil (→ caller falls back +// to base) for no-match / empty-capability / non-positive rate. +// +// Hermetic: exercises the pure resolver without the CGO ffmpeg toolchain, the +// remote-signer HTTP handler, or Kafka. +func TestResolveByocPrice(t *testing.T) { + require := require.New(t) + + byocCap := uint32(core.Capability_BYOC) + oInfo := &net.OrchestratorInfo{ + PriceInfo: &net.PriceInfo{PricePerUnit: 100, PixelsPerUnit: 1}, // base (unused by resolver) + CapabilitiesPrices: []*net.PriceInfo{ + {Capability: byocCap, Constraint: "nano-banana", PricePerUnit: 10, PixelsPerUnit: 1}, + {Capability: byocCap, Constraint: "recraft-v4", PricePerUnit: 20, PixelsPerUnit: 1}, + // decoy: a non-BYOC capability sharing the constraint name must NOT match + {Capability: uint32(core.Capability_LiveVideoToVideo), Constraint: "nano-banana", PricePerUnit: 999, PixelsPerUnit: 1}, + // decoy: zero/invalid BYOC rate must be treated as no-match (fallback) + {Capability: byocCap, Constraint: "free-cap", PricePerUnit: 0, PixelsPerUnit: 1}, + // misconfig: an early invalid duplicate must NOT shadow a later valid + // entry for the same constraint (scan continues past invalid rates). + {Capability: byocCap, Constraint: "dup-cap", PricePerUnit: 0, PixelsPerUnit: 1}, + {Capability: byocCap, Constraint: "dup-cap", PricePerUnit: 30, PixelsPerUnit: 1}, + }, + } + + tests := []struct { + name string + capability string + oInfo *net.OrchestratorInfo + want *net.PriceInfo + }{ + { + name: "resolves per capability", + capability: "recraft-v4", + oInfo: oInfo, + want: &net.PriceInfo{PricePerUnit: 20, PixelsPerUnit: 1}, + }, + { + name: "resolves the other capability", + capability: "nano-banana", + oInfo: oInfo, + want: &net.PriceInfo{PricePerUnit: 10, PixelsPerUnit: 1}, + }, + { + name: "unknown capability falls back (nil)", + capability: "does-not-exist", + oInfo: oInfo, + want: nil, + }, + { + name: "empty capability falls back (nil)", + capability: "", + oInfo: oInfo, + want: nil, + }, + { + name: "zero/invalid matched rate falls back (nil)", + capability: "free-cap", + oInfo: oInfo, + want: nil, + }, + { + name: "skips invalid duplicate, honors later valid entry", + capability: "dup-cap", + oInfo: oInfo, + want: &net.PriceInfo{PricePerUnit: 30, PixelsPerUnit: 1}, + }, + { + name: "no capabilities prices falls back (nil)", + capability: "nano-banana", + oInfo: &net.OrchestratorInfo{PriceInfo: &net.PriceInfo{PricePerUnit: 100, PixelsPerUnit: 1}}, + want: nil, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + req := &RemotePaymentRequest{Capability: tc.capability} + got := resolveByocPrice(req, tc.oInfo) + if tc.want == nil { + require.Nil(got) + return + } + require.NotNil(got) + require.Equal(tc.want.PricePerUnit, got.PricePerUnit, "PricePerUnit") + require.Equal(tc.want.PixelsPerUnit, got.PixelsPerUnit, "PixelsPerUnit") + }) + } +} + +// TestGenerateLivePayment_ByocPerCapPricing drives the full GenerateLivePayment +// handler (no CGO needed) and proves, via the resulting signed balance/state: +// - flag OFF: a matching cap is ignored; the fee is the byte-identical lv2v +// synthetic-pixel base fee (zero-regression); +// - flag ON: the fee is the resolved per-capability rate * compute-seconds; +// - flag ON: two caps in a 2:1 tariff produce a 2:1 fee ratio; +// - flag ON, unknown cap: falls back to the base lv2v fee; +// - flag ON with base >> 2x cap returns 200 (no false "price more than +// doubled"), proving the oInfo.PriceInfo = capPrice correction. +func TestGenerateLivePayment_ByocPerCapPricing(t *testing.T) { + require := require.New(t) + + ethClient := newTestEthClient(t) + node, _ := core.NewLivepeerNode(ethClient, "", nil) + node.Balances = core.NewAddressBalances(1 * time.Minute) + + // Large per-ticket EV keeps the lv2v fee within the handler's 100-ticket cap + // (lv2v fee = base * 1280*720*30*60 is large), while the small per-second + // BYOC fees resolve to a single ticket. + ev := big.NewRat(10_000_000_000, 1) + var totalTickets uint32 + sender := newMockSender(mockSenderConfig{ + ev: ev, + createTicketBatchFn: func(args mock.Arguments, batch *pm.TicketBatch) { + size := args.Int(1) + *batch = *defaultTicketBatch() + var baseSig []byte + if len(batch.SenderParams) > 0 && batch.SenderParams[0] != nil { + baseSig = batch.SenderParams[0].Sig + } + batch.SenderParams = make([]*pm.TicketSenderParams, size) + for i := 0; i < size; i++ { + totalTickets++ + batch.SenderParams[i] = &pm.TicketSenderParams{SenderNonce: totalTickets, Sig: baseSig} + } + }, + }) + node.Sender = sender + ls := &LivepeerServer{LivepeerNode: node} + + // Generous global max price so neither base nor cap prices trip the ceiling. + autoPrice, err := core.NewAutoConvertedPrice("", big.NewRat(1_000, 1), nil) + require.NoError(err) + BroadcastCfg.SetMaxPrice(autoPrice) + defer BroadcastCfg.SetMaxPrice(nil) + + // base is intentionally >> 2x the cap rates so a regression that left + // oInfo.PriceInfo at base (while locking the cap price into state) would trip + // the >2x doubling guard. With the correction it must NOT trip. + const basePPU, capNanoPPU, capRecraftPPU = 100, 10, 20 + oInfo := &net.OrchestratorInfo{ + Address: ethClient.addr.Bytes(), + PriceInfo: &net.PriceInfo{PricePerUnit: basePPU, PixelsPerUnit: 1}, + TicketParams: &net.TicketParams{ + Recipient: pm.RandAddress().Bytes(), + }, + AuthToken: stubAuthToken, + CapabilitiesPrices: []*net.PriceInfo{ + {Capability: uint32(core.Capability_BYOC), Constraint: "nano-banana", PricePerUnit: capNanoPPU, PixelsPerUnit: 1}, + {Capability: uint32(core.Capability_BYOC), Constraint: "recraft-v4", PricePerUnit: capRecraftPPU, PixelsPerUnit: 1}, + }, + } + orchBlob, err := proto.Marshal(oInfo) + require.NoError(err) + + // Fresh-session preload bills 60 seconds (lv2v) of data. + const preloadSecs = 60 + lv2vPixels := int64(defaultSegInfo.Height) * int64(defaultSegInfo.Width) * int64(defaultSegInfo.FPS) * preloadSecs + + doPayment := func(capability string) RemotePaymentState { + reqBody, err := json.Marshal(RemotePaymentRequest{ + Orchestrator: orchBlob, + Type: RemoteType_LiveVideoToVideo, + Capability: capability, + }) + require.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/generate-live-payment", bytes.NewReader(reqBody)) + rr := httptest.NewRecorder() + ls.GenerateLivePayment(rr, req) + require.Equal(http.StatusOK, rr.Code, "body=%s", rr.Body.String()) + var resp RemotePaymentResponse + require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) + var state RemotePaymentState + require.NoError(json.Unmarshal(resp.State.State, &state)) + return state + } + + // fee = numTickets*ev - balance for a fresh session; recover it from the + // signed state's balance and the locked initial price (also returned). + feeFromState := func(state RemotePaymentState) *big.Rat { + bal := new(big.Rat) + _, ok := bal.SetString(state.Balance) + require.True(ok, "parse balance %q", state.Balance) + // numTickets*ev = balance + fee, and numTickets*ev is the smallest + // multiple of ev that is >= max(fee, ev); recover fee = k*ev - balance + // where k = round((balance)/ev upward to the enclosing ticket). Simpler: + // derive fee from initialPrice * pixels directly using the locked state. + price := big.NewRat(state.InitialPricePerUnit, state.InitialPixelsPerUnit) + var pixels int64 + if state.InitialPricePerUnit == basePPU && state.InitialPixelsPerUnit == 1 { + pixels = lv2vPixels // base path: lv2v synthetic pixels + } else { + pixels = preloadSecs // byoc path: compute-seconds + } + return new(big.Rat).Mul(price, big.NewRat(pixels, 1)) + } + + assertBalanceConsistent := func(state RemotePaymentState, fee *big.Rat) { + // Independent check: balance must equal numTickets*ev - fee with + // numTickets = ceil(max(fee,ev)/ev), validating the fee end-to-end. + smc := fee + if ev.Cmp(smc) > 0 { + smc = ev + } + q := new(big.Rat).Quo(smc, ev) + nt := new(big.Int).Quo(q.Num(), q.Denom()) + if new(big.Int).Rem(q.Num(), q.Denom()).Sign() != 0 { + nt.Add(nt, big.NewInt(1)) + } + wantBal := new(big.Rat).Sub(new(big.Rat).Mul(new(big.Rat).SetInt(nt), ev), fee) + gotBal := new(big.Rat) + _, ok := gotBal.SetString(state.Balance) + require.True(ok) + require.Zero(gotBal.Cmp(wantBal), "balance got=%s want=%s fee=%s", gotBal.RatString(), wantBal.RatString(), fee.RatString()) + } + + // --- flag OFF: cap present but ignored → base lv2v synthetic-pixel fee --- + ls.LivepeerNode.ByocPerCapPricing = false + offState := doPayment("nano-banana") + require.EqualValues(basePPU, offState.InitialPricePerUnit, "flag OFF must lock the base price") + require.EqualValues(1, offState.InitialPixelsPerUnit) + offFee := new(big.Rat).Mul(big.NewRat(basePPU, 1), big.NewRat(lv2vPixels, 1)) + require.Zero(feeFromState(offState).Cmp(offFee)) + assertBalanceConsistent(offState, offFee) + + // --- flag ON --- + ls.LivepeerNode.ByocPerCapPricing = true + + // nano-banana: fee = capNanoPPU * 60 (seconds basis); base>>2x cap must NOT + // trip the doubling guard (proves oInfo.PriceInfo = capPrice). + nanoState := doPayment("nano-banana") + require.EqualValues(capNanoPPU, nanoState.InitialPricePerUnit, "flag ON must lock the resolved cap price") + require.EqualValues(1, nanoState.InitialPixelsPerUnit) + nanoFee := big.NewRat(capNanoPPU*preloadSecs, 1) + require.Zero(feeFromState(nanoState).Cmp(nanoFee), "nano fee") + assertBalanceConsistent(nanoState, nanoFee) + + // recraft-v4 priced 2x nano → fee ratio is exactly 2:1. + recraftState := doPayment("recraft-v4") + require.EqualValues(capRecraftPPU, recraftState.InitialPricePerUnit) + recraftFee := big.NewRat(capRecraftPPU*preloadSecs, 1) + require.Zero(feeFromState(recraftState).Cmp(recraftFee), "recraft fee") + assertBalanceConsistent(recraftState, recraftFee) + require.Zero(recraftFee.Cmp(new(big.Rat).Mul(nanoFee, big.NewRat(2, 1))), "recraft fee must be 2x nano fee") + + // unknown cap with flag ON → fall back to base lv2v fee (zero-regression). + unknownState := doPayment("totally-unknown-cap") + require.EqualValues(basePPU, unknownState.InitialPricePerUnit, "unknown cap must fall back to base") + require.Zero(feeFromState(unknownState).Cmp(offFee)) + assertBalanceConsistent(unknownState, offFee) + + // type:"byoc" + capabilities protobuf (no capability string) must bill like BYOC. + doPaymentByocType := func(modelID string) RemotePaymentState { + reqBody, err := json.Marshal(RemotePaymentRequest{ + Orchestrator: orchBlob, + Type: RemoteType_BYOC, + Capabilities: byocCapsProtoBytes(t, modelID), + }) + require.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/generate-live-payment", bytes.NewReader(reqBody)) + rr := httptest.NewRecorder() + ls.GenerateLivePayment(rr, req) + require.Equal(http.StatusOK, rr.Code, "body=%s", rr.Body.String()) + var resp RemotePaymentResponse + require.NoError(json.NewDecoder(rr.Body).Decode(&resp)) + var state RemotePaymentState + require.NoError(json.Unmarshal(resp.State.State, &state)) + return state + } + + byocTypeState := doPaymentByocType("nano-banana") + require.EqualValues(capNanoPPU, byocTypeState.InitialPricePerUnit, "type byoc must resolve per-cap price from capabilities proto") + require.Zero(feeFromState(byocTypeState).Cmp(nanoFee), "type byoc fee") +} diff --git a/server/rpc.go b/server/rpc.go index da335b98aa..471e5b5fa5 100644 --- a/server/rpc.go +++ b/server/rpc.go @@ -13,8 +13,11 @@ import ( "sync" "time" + "golang.org/x/net/http2" + "golang.org/x/net/http2/h2c" "google.golang.org/grpc" "google.golang.org/grpc/credentials" + "google.golang.org/grpc/credentials/insecure" "github.com/livepeer/go-livepeer/ai/worker" "github.com/livepeer/go-livepeer/byoc" @@ -50,6 +53,7 @@ type Orchestrator interface { ServiceURI() *url.URL Address() ethcommon.Address TranscoderSecret() string + RegistrationSecret() string Sign([]byte) ([]byte, error) VerifySig(ethcommon.Address, string, []byte) bool CheckCapacity(core.ManifestID) error @@ -161,6 +165,11 @@ type GetOrchestratorInfoParams struct { IgnoreCapacityCheck bool } +type refreshPaymentParamsRequest struct { + Sender string `json:"sender"` + ManifestID string `json:"manifest_id"` +} + func (bs *BroadcastSession) Transcoder() string { bs.lock.RLock() defer bs.lock.RUnlock() @@ -238,6 +247,7 @@ func StartTranscodeServer(orch Orchestrator, bind string, mux *http.ServeMux, wo net.RegisterOrchestratorServer(s, &lp) lp.transRPC.HandleFunc("/segment", lp.ServeSegment) lp.transRPC.HandleFunc("/payment", lp.Payment) + lp.transRPC.HandleFunc("POST /refresh-payment", lp.RefreshPayment) if acceptRemoteTranscoders { net.RegisterTranscoderServer(s, &lp) lp.transRPC.HandleFunc("/transcodeResults", lp.TranscodeResults) @@ -254,23 +264,56 @@ func StartTranscodeServer(orch Orchestrator, bind string, mux *http.ServeMux, wo //API for dynamic capabilities lp.byocSrv = byoc.NewBYOCOrchestratorServer(n, orch, lp.trickleSrv, TrickleHTTPPath, lp.transRPC) - cert, key, err := getCert(orch.ServiceURI(), workDir) + stopTrickle := lp.trickleSrv.Start() + defer stopTrickle() + + bind, listenerTLS, err := parseHTTPAddr(bind) if err != nil { return err } - stopTrickle := lp.trickleSrv.Start() - defer stopTrickle() - - glog.Info("Listening for RPC on ", bind) + scheme := "https" + handler := http.Handler(&lp) + if !listenerTLS { + scheme = "http" + handler = h2c.NewHandler(&lp, &http2.Server{}) + } + glog.Infof("Listening for RPC on %s://%s", scheme, bind) srv := http.Server{ Addr: bind, - Handler: &lp, + Handler: handler, IdleTimeout: HTTPIdleTimeout, } + if !listenerTLS { + return srv.ListenAndServe() + } + + cert, key, err := getCert(orch.ServiceURI(), workDir) + if err != nil { + return err + } return srv.ListenAndServeTLS(cert, key) } +func parseHTTPAddr(addr string) (string, bool, error) { + if !strings.Contains(addr, "://") { + return addr, true, nil + } + + uri, err := url.Parse(addr) + if err != nil { + return "", false, err + } + if uri.Scheme != "http" && uri.Scheme != "https" { + return "", false, fmt.Errorf("unsupported -httpAddr scheme %q", uri.Scheme) + } + if uri.Path != "" && uri.Path != "/" { + return "", false, fmt.Errorf("-httpAddr must not include a path") + } + + return uri.Host, uri.Scheme != "http", nil +} + // CheckOrchestratorAvailability - the broadcaster calls CheckOrchestratorAvailability which invokes Ping on the orchestrator func CheckOrchestratorAvailability(orch Orchestrator) bool { ts := time.Now() @@ -299,6 +342,38 @@ func CheckOrchestratorAvailability(orch Orchestrator) bool { return orch.VerifySig(orch.Address(), string(ping), pong.Value) } +func CheckOrchestratorDiscoveryAvailability(orch Orchestrator) bool { + uri := orch.ServiceURI().JoinPath("discovery") + ctx, cancel := context.WithTimeout(context.Background(), GRPCTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, uri.String(), nil) + if err != nil { + glog.Error("Unable to create discovery availability request: ", err) + return false + } + + resp, err := httpClient.Do(req) + if err != nil { + glog.Error("Was not able to submit discovery availability request: ", err) + return false + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + glog.Errorf("Discovery availability check failed with status %d", resp.StatusCode) + return false + } + + var discovery []json.RawMessage + if err := json.NewDecoder(resp.Body).Decode(&discovery); err != nil { + glog.Error("Unable to decode discovery availability response: ", err) + return false + } + + return true +} + func ping(context context.Context, req *net.PingPong, orch Orchestrator) (*net.PingPong, error) { glog.Info("Received Ping request") value, err := orch.Sign(req.Value) @@ -348,8 +423,12 @@ func EndTranscodingSession(ctx context.Context, sess *BroadcastSession) error { func startOrchestratorClient(ctx context.Context, uri *url.URL) (net.OrchestratorClient, *grpc.ClientConn, error) { clog.V(common.DEBUG).Infof(ctx, "Connecting RPC to uri=%v", uri) + transportCredentials := credentials.NewTLS(tlsConfig) + if uri.Scheme == "http" { + transportCredentials = insecure.NewCredentials() + } conn, err := grpc.Dial(uri.Host, - grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), + grpc.WithTransportCredentials(transportCredentials), grpc.WithBlock(), grpc.WithTimeout(GRPCConnectTimeout)) if err != nil { @@ -369,6 +448,78 @@ func genEndSessionRequest(sess *BroadcastSession) (*net.EndTranscodingSessionReq return &net.EndTranscodingSessionRequest{AuthToken: sess.OrchestratorInfo.AuthToken}, nil } +func (h *lphttp) RefreshPayment(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + var req refreshPaymentParamsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + respondJsonError(ctx, w, err, http.StatusBadRequest) + return + } + + if req.Sender == "" || !ethcommon.IsHexAddress(req.Sender) { + respondJsonError(ctx, w, fmt.Errorf("invalid sender"), http.StatusBadRequest) + return + } + if strings.TrimSpace(req.ManifestID) == "" { + respondJsonError(ctx, w, fmt.Errorf("missing manifest_id"), http.StatusBadRequest) + return + } + + sender := ethcommon.HexToAddress(req.Sender) + manifestID := core.ManifestID(strings.TrimSpace(req.ManifestID)) + priceInfo, ok, err := fixedPriceInfo(h.node, sender, manifestID) + if err != nil { + respondJsonError(ctx, w, err, http.StatusInternalServerError) + return + } + if !ok { + respondJsonError(ctx, w, fmt.Errorf("fixed price not found for session"), http.StatusConflict) + return + } + + ticketParams, err := h.orchestrator.TicketParams(sender, priceInfo) + if err != nil { + respondJsonError(ctx, w, err, http.StatusInternalServerError) + return + } + + expiration := time.Now().Add(authTokenValidPeriod).Unix() + oInfo := &net.OrchestratorInfo{ + Transcoder: h.orchestrator.ServiceURI().String(), + TicketParams: ticketParams, + PriceInfo: priceInfo, + Address: h.orchestrator.Address().Bytes(), + AuthToken: h.orchestrator.AuthToken(string(manifestID), expiration), + } + + w.Header().Set("Content-Type", "application/json") + data, err := marshalLivePaymentChallengeResponse(oInfo) + if err != nil { + respondJsonError(ctx, w, err, http.StatusInternalServerError) + return + } + if _, err := w.Write(data); err != nil { + clog.Errorf(ctx, "Error encoding refreshed payment params err=%v", err) + } +} + +func fixedPriceInfo(node *core.LivepeerNode, sender ethcommon.Address, manifestID core.ManifestID) (*net.PriceInfo, bool, error) { + if node == nil || node.Balances == nil { + return nil, false, nil + } + fixedPrice := node.Balances.FixedPrice(sender, manifestID) + if fixedPrice == nil { + return nil, false, nil + } + if !fixedPrice.Num().IsInt64() || !fixedPrice.Denom().IsInt64() { + return nil, true, fmt.Errorf("fixed price cannot be represented as int64 price info") + } + return &net.PriceInfo{ + PricePerUnit: fixedPrice.Num().Int64(), + PixelsPerUnit: fixedPrice.Denom().Int64(), + }, true, nil +} + func getOrchestrator(orch Orchestrator, req *net.OrchestratorRequest) (*net.OrchestratorInfo, error) { addr := ethcommon.BytesToAddress(req.Address) if err := verifyOrchestratorReq(orch, addr, req.Sig); err != nil { diff --git a/server/rpc_test.go b/server/rpc_test.go index fd81e05acb..bce71a09f8 100644 --- a/server/rpc_test.go +++ b/server/rpc_test.go @@ -8,15 +8,21 @@ import ( "errors" "fmt" "math/big" + gonet "net" "net/http" "net/http/httptest" "net/url" + "strings" + "sync" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" + "golang.org/x/net/http2" + "golang.org/x/net/http2/h2c" + "google.golang.org/grpc" "pgregory.net/rapid" "github.com/ethereum/go-ethereum/accounts" @@ -69,18 +75,25 @@ func (m *mockBalance) Clear() { } type stubOrchestrator struct { - priv *ecdsa.PrivateKey - block *big.Int - signErr error - sessCapErr error - ticketParams *net.TicketParams - priceInfo *net.PriceInfo - serviceURI string - res *core.TranscodeResult - offchain bool - caps *core.Capabilities - authToken *net.AuthToken - jobPriceInfo *net.PriceInfo + priv *ecdsa.PrivateKey + block *big.Int + signErr error + sessCapErr error + ticketParams *net.TicketParams + priceInfo *net.PriceInfo + serviceURI string + res *core.TranscodeResult + offchain bool + caps *core.Capabilities + authToken *net.AuthToken + jobPriceInfo *net.PriceInfo + secret string + balanceMu sync.Mutex + balances map[ethcommon.Address]map[core.ManifestID]*big.Rat + paymentCredit *big.Rat + requestMu sync.Mutex + storageReqs []string + lv2vReqs []string } func (r *stubOrchestrator) GetLiveAICapacity(pipeline, modelID string) worker.Capacity { @@ -95,6 +108,62 @@ func (r *stubOrchestrator) ServiceURI() *url.URL { return url } +func (r *stubOrchestrator) LiveRunnerURI() *url.URL { + return r.ServiceURI() +} + +func TestStartOrchestratorClientHTTP(t *testing.T) { + listener, err := gonet.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + orchRPC := grpc.NewServer() + orch := &stubOrchestrator{ + offchain: true, + serviceURI: "http://" + listener.Addr().String(), + } + lp := &lphttp{ + orchestrator: orch, + orchRPC: orchRPC, + transRPC: http.NewServeMux(), + } + net.RegisterOrchestratorServer(orchRPC, lp) + + srv := &http.Server{ + Handler: h2c.NewHandler(lp, &http2.Server{}), + } + go func() { + _ = srv.Serve(listener) + }() + defer srv.Shutdown(context.Background()) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + client, conn, err := startOrchestratorClient(ctx, orch.ServiceURI()) + require.NoError(t, err) + defer conn.Close() + + _, err = client.Ping(ctx, &net.PingPong{Value: []byte("ping")}) + require.NoError(t, err) +} + +func TestParseHTTPAddrScheme(t *testing.T) { + addr, listenerTLS, err := parseHTTPAddr("http://:8935") + require.NoError(t, err) + require.Equal(t, ":8935", addr) + require.False(t, listenerTLS) + + addr, listenerTLS, err = parseHTTPAddr("https://127.0.0.1:8935") + require.NoError(t, err) + require.Equal(t, "127.0.0.1:8935", addr) + require.True(t, listenerTLS) + + addr, listenerTLS, err = parseHTTPAddr("127.0.0.1:8935") + require.NoError(t, err) + require.Equal(t, "127.0.0.1:8935", addr) + require.True(t, listenerTLS) +} + func (r *stubOrchestrator) Nodes() []string { return nil } @@ -144,6 +213,18 @@ func (r *stubOrchestrator) StreamIDs(jobID string) ([]core.StreamID, error) { } func (r *stubOrchestrator) ProcessPayment(ctx context.Context, payment net.Payment, manifestID core.ManifestID) error { + if r.balances != nil && r.paymentCredit != nil { + r.balanceMu.Lock() + defer r.balanceMu.Unlock() + sender := getPaymentSender(payment) + if r.balances[sender] == nil { + r.balances[sender] = make(map[core.ManifestID]*big.Rat) + } + if r.balances[sender][manifestID] == nil { + r.balances[sender][manifestID] = big.NewRat(0, 1) + } + r.balances[sender][manifestID].Add(r.balances[sender][manifestID], r.paymentCredit) + } return nil } @@ -160,13 +241,39 @@ func (r *stubOrchestrator) GetCapabilitiesPrices(sender ethcommon.Address) ([]*n } func (r *stubOrchestrator) SufficientBalance(addr ethcommon.Address, manifestID core.ManifestID) bool { + if r.balances != nil { + balance := r.Balance(addr, manifestID) + return balance != nil && balance.Sign() > 0 + } return true } func (r *stubOrchestrator) DebitFees(addr ethcommon.Address, manifestID core.ManifestID, price *net.PriceInfo, pixels int64) { + if r.balances == nil { + return + } + priceRat := big.NewRat(price.GetPricePerUnit(), price.GetPixelsPerUnit()) + fee := priceRat.Mul(priceRat, big.NewRat(pixels, 1)) + r.balanceMu.Lock() + defer r.balanceMu.Unlock() + if r.balances[addr] == nil { + r.balances[addr] = make(map[core.ManifestID]*big.Rat) + } + if r.balances[addr][manifestID] == nil { + r.balances[addr][manifestID] = big.NewRat(0, 1) + } + r.balances[addr][manifestID].Sub(r.balances[addr][manifestID], fee) } func (r *stubOrchestrator) Balance(addr ethcommon.Address, manifestID core.ManifestID) *big.Rat { + if r.balances != nil { + r.balanceMu.Lock() + defer r.balanceMu.Unlock() + if r.balances[addr] == nil || r.balances[addr][manifestID] == nil { + return nil + } + return new(big.Rat).Set(r.balances[addr][manifestID]) + } return big.NewRat(0, 1) } @@ -207,7 +314,10 @@ func (r *stubOrchestrator) ServeTranscoder(stream net.Transcoder_RegisterTransco func (r *stubOrchestrator) TranscoderResults(job int64, res *core.RemoteTranscoderResult) { } func (r *stubOrchestrator) TranscoderSecret() string { - return "" + return r.secret +} +func (r *stubOrchestrator) RegistrationSecret() string { + return r.TranscoderSecret() } func (r *stubOrchestrator) PriceInfoForCaps(sender ethcommon.Address, manifestID core.ManifestID, caps *net.Capabilities) (*net.PriceInfo, error) { return &net.PriceInfo{PricePerUnit: 4, PixelsPerUnit: 1}, nil @@ -241,6 +351,9 @@ func (r *stubOrchestrator) TextToSpeech(ctx context.Context, requestID string, r } func (r *stubOrchestrator) LiveVideoToVideo(ctx context.Context, requestID string, req worker.GenLiveVideoToVideoJSONRequestBody) (interface{}, error) { + r.requestMu.Lock() + defer r.requestMu.Unlock() + r.lv2vReqs = append(r.lv2vReqs, requestID) return nil, nil } @@ -250,6 +363,9 @@ func (r *stubOrchestrator) CheckAICapacity(pipeline, modelID string) (bool, chan func (r *stubOrchestrator) AIResults(job int64, res *core.RemoteAIWorkerResult) { } func (r *stubOrchestrator) CreateStorageForRequest(requestID string) error { + r.requestMu.Lock() + defer r.requestMu.Unlock() + r.storageReqs = append(r.storageReqs, requestID) return nil } func (r *stubOrchestrator) GetStorageForRequest(requestID string) (drivers.OSSession, bool) { @@ -758,6 +874,81 @@ func TestPing(t *testing.T) { } } +func TestCheckOrchestratorDiscoveryAvailability(t *testing.T) { + tests := []struct { + name string + statusCode int + body string + want bool + }{ + { + name: "empty discovery succeeds", + statusCode: http.StatusOK, + body: `[]`, + want: true, + }, + { + name: "nonempty discovery succeeds", + statusCode: http.StatusOK, + body: `[{"address":"http://localhost:1234","runners":[]}]`, + want: true, + }, + { + name: "not found fails", + statusCode: http.StatusNotFound, + body: `live runners are not supported`, + want: false, + }, + { + name: "invalid json fails", + statusCode: http.StatusOK, + body: `not-json`, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodGet, r.Method) + require.Equal(t, "/discovery", r.URL.Path) + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.body)) + })) + defer srv.Close() + + orch := newStubOrchestrator() + orch.serviceURI = srv.URL + require.Equal(t, tt.want, CheckOrchestratorDiscoveryAvailability(orch)) + }) + } +} + +func TestCheckOrchestratorDiscoveryAvailabilityConnectionError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + uri := srv.URL + srv.Close() + + orch := newStubOrchestrator() + orch.serviceURI = uri + require.False(t, CheckOrchestratorDiscoveryAvailability(orch)) +} + +func TestCheckOrchestratorDiscoveryAvailabilityTLS(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/discovery", r.URL.Path) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`[]`)) + })) + defer srv.Close() + + orch := newStubOrchestrator() + orch.serviceURI = srv.URL + require.True(t, CheckOrchestratorDiscoveryAvailability(orch)) +} + func TestValidatePrice(t *testing.T) { assert := assert.New(t) mid := core.RandomManifestID() @@ -1136,6 +1327,90 @@ func TestGetOrchestrator_StorageInit(t *testing.T) { drivers.NodeStorage = drivers.NewMemoryDriver(nil) } +func TestRefreshPayment_ReturnsPinnedPaymentChallenge(t *testing.T) { + sender := ethcommon.HexToAddress("0x1234567890123456789012345678901234567890") + manifestID := core.ManifestID("manifest-id") + fixedPrice := big.NewRat(7, 3) + ticketParams := defaultTicketParams() + authToken := &net.AuthToken{Token: []byte("token"), SessionId: string(manifestID), Expiration: time.Now().Add(time.Hour).Unix()} + orchAddr := ethcommon.HexToAddress("0xabcdefabcdefabcdefabcdefabcdefabcdefabcd") + + node := &core.LivepeerNode{Balances: core.NewAddressBalances(time.Hour)} + node.Balances.Credit(sender, manifestID, big.NewRat(0, 1)) + node.Balances.SetFixedPrice(sender, manifestID, fixedPrice) + + orch := &mockOrchestrator{} + orch.On("TicketParams", sender, mock.MatchedBy(func(price *net.PriceInfo) bool { + return price != nil && price.PricePerUnit == 7 && price.PixelsPerUnit == 3 + })).Return(ticketParams, nil) + orch.On("ServiceURI").Return(mustParseUrl(t, "http://orch.example")) + orch.On("Address").Return(orchAddr) + orch.On("AuthToken", string(manifestID), mock.Anything).Return(authToken) + + lp := &lphttp{orchestrator: orch, node: node} + body := `{"sender":"` + sender.Hex() + `","manifest_id":"` + string(manifestID) + `"}` + req := httptest.NewRequest(http.MethodPost, "/refresh-payment", strings.NewReader(body)) + rr := httptest.NewRecorder() + + lp.RefreshPayment(rr, req) + + require.Equal(t, http.StatusOK, rr.Code) + challenge, got := decodeLiveRunnerPaymentChallenge(t, rr.Body.Bytes()) + require.Equal(t, "http://orch.example", challenge.Orchestrator) + require.Equal(t, string(manifestID), challenge.ManifestID) + require.Equal(t, "http://orch.example", got.Transcoder) + require.True(t, proto.Equal(ticketParams, got.TicketParams)) + require.Equal(t, int64(7), got.PriceInfo.PricePerUnit) + require.Equal(t, int64(3), got.PriceInfo.PixelsPerUnit) + require.Equal(t, orchAddr.Bytes(), got.Address) + require.True(t, proto.Equal(authToken, got.AuthToken)) + require.Nil(t, got.Capabilities) + require.Nil(t, got.Storage) + require.Nil(t, got.Hardware) + require.Nil(t, got.CapabilitiesPrices) + require.Nil(t, got.Nodes) + orch.AssertNotCalled(t, "VerifySig", mock.Anything, mock.Anything, mock.Anything) +} + +func TestRefreshPayment_NoFixedPrice_ReturnsConflict(t *testing.T) { + node := &core.LivepeerNode{Balances: core.NewAddressBalances(time.Hour)} + lp := &lphttp{node: node} + req := httptest.NewRequest( + http.MethodPost, + "/refresh-payment", + strings.NewReader(`{"sender":"0x1234567890123456789012345678901234567890","manifest_id":"manifest-id"}`), + ) + rr := httptest.NewRecorder() + + lp.RefreshPayment(rr, req) + + require.Equal(t, http.StatusConflict, rr.Code) + require.Contains(t, rr.Body.String(), "fixed price not found for session") +} + +func TestRefreshPayment_InvalidRequest(t *testing.T) { + lp := &lphttp{} + tests := []struct { + name string + body string + }{ + {name: "missing sender", body: `{"manifest_id":"manifest-id"}`}, + {name: "invalid sender", body: `{"sender":"not-an-address","manifest_id":"manifest-id"}`}, + {name: "missing manifest id", body: `{"sender":"0x1234567890123456789012345678901234567890"}`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/refresh-payment", strings.NewReader(tt.body)) + rr := httptest.NewRecorder() + + lp.RefreshPayment(rr, req) + + require.Equal(t, http.StatusBadRequest, rr.Code) + }) + } +} + func TestGetPriceInfo_NoWebhook_DefaultPriceError_ReturnsError(t *testing.T) { assert := assert.New(t) orch := &mockOrchestrator{} @@ -1483,6 +1758,9 @@ func (o *mockOrchestrator) TranscoderSecret() string { o.Called() return "" } +func (o *mockOrchestrator) RegistrationSecret() string { + return o.TranscoderSecret() +} func (o *mockOrchestrator) Sign(msg []byte) ([]byte, error) { o.Called(msg) return nil, nil diff --git a/server/webserver.go b/server/webserver.go index d05e0c399f..65764b9b4e 100644 --- a/server/webserver.go +++ b/server/webserver.go @@ -1,13 +1,16 @@ package server import ( + "encoding/json" "flag" + "io" "net/http" // pprof adds handlers to default mux via `init()` _ "net/http/pprof" "github.com/golang/glog" + "github.com/livepeer/go-livepeer/ai/runner" "github.com/livepeer/go-livepeer/monitor" ) @@ -52,6 +55,7 @@ func (s *LivepeerServer) cliWebServerHandlers(bindAddr string) *http.ServeMux { mux.Handle("/setMaxPriceForCapability", mustHaveFormParams(s.setMaxPriceForCapability(), "maxPricePerUnit", "pixelsPerUnit", "currency", "pipeline", "modelID")) mux.Handle("/getAISessionPoolsInfo", s.getAIPoolsInfoHandler()) mux.Handle("/getNetworkCapabilities", s.getNetworkCapabilitiesHandler()) + mux.Handle("/registerLiveRunners", s.registerLiveRunnersHandler()) // Rounds mux.Handle("/currentRound", currentRoundHandler(client)) @@ -119,3 +123,41 @@ func (s *LivepeerServer) cliWebServerHandlers(bindAddr string) *http.ServeMux { return mux } + +func (s *LivepeerServer) registerLiveRunnersHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + manager, ok := s.LivepeerNode.LiveRunnerManager.(interface { + RegisterStaticRunnersJSON([]byte) (*runner.StaticLiveRunnerRegistrationResponse, error) + }) + if !ok { + http.Error(w, "live runners are not supported", http.StatusNotFound) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp, err := manager.RegisterStaticRunnersJSON(body) + if err != nil { + statusCode := http.StatusBadRequest + if runnerErr, ok := err.(*runner.RunnerError); ok { + statusCode = runnerErr.StatusCode + } + http.Error(w, err.Error(), statusCode) + return + } + data, err := json.Marshal(resp) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(data) + }) +} diff --git a/trickle/README.md b/trickle/README.md index ce4a8706b5..bd2cfd69e3 100644 --- a/trickle/README.md +++ b/trickle/README.md @@ -50,6 +50,11 @@ Servers may offer some grace with leading sequence numbers to avoid data races, Publishers are responsible for segmenting content (if necessary) and subscribers are responsible for re-assembling content (if necessary) +Successful publisher POST responses include `Lp-Trickle-Seq` metadata (HTTP header) with the effective segment index written by the server. + +If a publisher sends `Lp-Trickle-Reset`, the server treats it as a restart signal for any `seq` value. +The server closes prior segments to unblock waiting subscribers while still allowing preconnected/ahead segments. + Subscribers can initiate a subscribe with a `seq` of -1 to retrieve the most recent publish. With preconnects, the subscriber may be waiting for the *next* publish. For video this allows clients to eg, start streaming at the live edge of the next GOP. Subscribers can retrieve the current `seq` with the `Lp-Trickle-Seq` metadata (HTTP header). This is useful in case `-1` was used to initiate the subscription; the subscribing client can then pre-connect to `Lp-Trickle-Seq + 1` @@ -58,6 +63,9 @@ GET `/channel-name/next` returns the next segment seq as plain text in the respo If the channel does not exist, the server returns 404. If the channel is closed, the server also includes `Lp-Trickle-Closed: terminated`. +`Lp-Trickle-Seq` from a publisher write using `seq=-1` is not sufficient by itself to drive real-time ordering if writes overlap. +It is still useful for post-facto mapping/observability, reconciliation after segment close, and debugging dropped/misordered assumptions. + Subscribers can initiate a subscribe with a `seq` of -N to get the Nth-from-last segment. (TODO) The server should send subscribers `Lp-Trickle-Size` metadata to indicate the size of the content up until now. This allows clients to know where the live edge is, eg video implementations can decode-and-discard frames up until the edge to achieve immediate playback without waiting for the next segment. (TODO) diff --git a/trickle/local_publisher.go b/trickle/local_publisher.go index bc0082dd8e..9fe9aae481 100644 --- a/trickle/local_publisher.go +++ b/trickle/local_publisher.go @@ -5,6 +5,7 @@ import ( "io" "log/slog" "sync" + "time" ) // local (in-memory) publisher for trickle protocol @@ -32,7 +33,10 @@ func (c *TrickleLocalPublisher) CreateChannel() { } func (c *TrickleLocalPublisher) Write(data io.Reader) error { - stream := c.server.getOrCreateStream(c.channelName, c.mimeType, true) + stream := c.server.getOrCreateStream(c.channelName, c.mimeType, false) + if stream == nil { + return StreamNotFoundErr + } c.mu.Lock() seq := c.seq segment, exists := stream.getForWrite(seq) @@ -56,6 +60,12 @@ func (c *TrickleLocalPublisher) Write(data io.Reader) error { for { n, err := data.Read(buf) if n > 0 { + if totalRead == 0 { + stream.mutex.Lock() + stream.nextWrite = seq + 1 + stream.writeTime = time.Now() + stream.mutex.Unlock() + } segment.writeData(buf[:n]) totalRead += n } diff --git a/trickle/segment_buffer.go b/trickle/segment_buffer.go new file mode 100644 index 0000000000..896b1656b2 --- /dev/null +++ b/trickle/segment_buffer.go @@ -0,0 +1,146 @@ +package trickle + +import ( + "sync/atomic" +) + +const ( + segmentBufferInitialPageSize = 32 * 1024 + segmentBufferMaxPageSize = 1024 * 1024 +) + +type segmentPage struct { + start int + data []byte + written int +} + +type segmentPageList struct { + pages []*segmentPage +} + +// segmentBuffer is an append-only paged buffer tailored for Segment fanout. +// Pages are immutable containers whose published prefix is tracked atomically, +// allowing readers to take a lock-free fast path when data is already available. +type segmentBuffer struct { + pageList atomic.Pointer[segmentPageList] + published atomic.Int64 + totalWritten int + initialCap int + maxPageCap int + nextPageCap int +} + +func newSegmentBuffer() *segmentBuffer { + return newSegmentBufferWithPageCaps(segmentBufferInitialPageSize, segmentBufferMaxPageSize) +} + +func newSegmentBufferWithPageCaps(initialCap, maxPageCap int) *segmentBuffer { + if initialCap <= 0 { + initialCap = segmentBufferInitialPageSize + } + if maxPageCap < initialCap { + maxPageCap = initialCap + } + b := &segmentBuffer{ + initialCap: initialCap, + maxPageCap: maxPageCap, + nextPageCap: initialCap, + } + b.pageList.Store(&segmentPageList{pages: nil}) + return b +} + +func (b *segmentBuffer) write(data []byte) { + for len(data) > 0 { + page := b.ensureTailPage() + written := page.written + remaining := cap(page.data) - written + if remaining > len(data) { + remaining = len(data) + } + + end := written + remaining + copy(page.data[written:end], data[:remaining]) + page.written = end + + data = data[remaining:] + b.totalWritten += remaining + } + b.published.Store(int64(b.totalWritten)) +} + +func (b *segmentBuffer) ensureTailPage() *segmentPage { + list := b.pageList.Load() + if list != nil && len(list.pages) > 0 { + tail := list.pages[len(list.pages)-1] + if tail.written < cap(tail.data) { + return tail + } + } + + pageCap := b.nextPageCap + if pageCap == 0 { + pageCap = b.initialCap + } + page := &segmentPage{ + start: b.totalWritten, + data: make([]byte, pageCap), + } + newPages := make([]*segmentPage, 0, len(list.pages)+1) + newPages = append(newPages, list.pages...) + newPages = append(newPages, page) + b.pageList.Store(&segmentPageList{pages: newPages}) + + if b.nextPageCap < b.maxPageCap { + b.nextPageCap *= 2 + if b.nextPageCap > b.maxPageCap { + b.nextPageCap = b.maxPageCap + } + } + return page +} + +func (b *segmentBuffer) isEmpty() bool { + return b.published.Load() == 0 +} + +func (b *segmentBuffer) readChunk(pos int) ([]byte, int, bool, bool, int) { + if pos < 0 { + return nil, pos, false, true, int(b.published.Load()) + } + published := int(b.published.Load()) + if pos > published { + return nil, pos, false, true, published + } + if pos == published { + return nil, pos, false, false, published + } + + list := b.pageList.Load() + if list == nil { + return nil, pos, false, false, published + } + + for i, page := range list.pages { + pageEnd := published + if i+1 < len(list.pages) { + pageEnd = list.pages[i+1].start + if pageEnd > published { + pageEnd = published + } + } + if pos < page.start || pos >= pageEnd { + continue + } + + offset := pos - page.start + data := page.data[offset : pageEnd-page.start] + return data, pageEnd, pageEnd == published, false, published + } + + // We may have observed a newer published cursor than the currently visible + // page-list snapshot. Treat this as "not available yet" so callers can retry + // under synchronization instead of incorrectly signaling EOF. + return nil, pos, false, false, published +} diff --git a/trickle/segment_buffer_test.go b/trickle/segment_buffer_test.go new file mode 100644 index 0000000000..b733611c47 --- /dev/null +++ b/trickle/segment_buffer_test.go @@ -0,0 +1,68 @@ +package trickle + +import ( + "bytes" + "testing" +) + +func TestSegmentBufferReadChunkAcrossPages(t *testing.T) { + buf := newSegmentBuffer() + chunkA := bytes.Repeat([]byte("a"), segmentBufferInitialPageSize) + chunkB := bytes.Repeat([]byte("b"), segmentBufferInitialPageSize*2) + want := append(append([]byte{}, chunkA...), chunkB...) + + buf.write(chunkA) + buf.write(chunkB) + + var got []byte + nextPos := 0 + for { + data, next, atTail, invalid, _ := buf.readChunk(nextPos) + if invalid { + t.Fatalf("invalid cursor at pos=%d", nextPos) + } + if len(data) == 0 { + break + } + got = append(got, data...) + nextPos = next + if atTail { + break + } + } + if !bytes.Equal(got, want) { + t.Fatalf("buffer mismatch: got=%d want=%d", len(got), len(want)) + } + if nextPos != len(want) { + t.Fatalf("unexpected tail cursor: pos=%d", nextPos) + } +} + +func TestSegmentBufferTailGrowth(t *testing.T) { + buf := newSegmentBuffer() + chunkA := []byte("hello") + chunkB := []byte(" world") + + buf.write(chunkA) + + data, nextPos, atTail, invalid, _ := buf.readChunk(0) + if len(data) == 0 || !atTail || invalid { + t.Fatalf("initial read state len=%d atTail=%v invalid=%v", len(data), atTail, invalid) + } + if string(data) != "hello" { + t.Fatalf("unexpected initial data %q", string(data)) + } + + buf.write(chunkB) + + data, nextPos, atTail, invalid, _ = buf.readChunk(nextPos) + if len(data) == 0 || !atTail || invalid { + t.Fatalf("growth read state len=%d atTail=%v invalid=%v", len(data), atTail, invalid) + } + if string(data) != " world" { + t.Fatalf("unexpected growth data %q", string(data)) + } + if nextPos != len(chunkA)+len(chunkB) { + t.Fatalf("unexpected grown cursor: pos=%d", nextPos) + } +} diff --git a/trickle/trickle_server.go b/trickle/trickle_server.go index 895677f9cf..b08a3a858c 100644 --- a/trickle/trickle_server.go +++ b/trickle/trickle_server.go @@ -13,6 +13,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" ) @@ -36,6 +37,14 @@ type TrickleServerConfig struct { // How often to sweep for idle channels (default 1 minute) SweepInterval time.Duration + + // BeforeCreate runs before HTTP channel creation. + // Return RequestError for expected client/policy failures. + BeforeCreate func(r *http.Request, streamName string) error + + // BeforeDelete runs before HTTP channel deletion. + // Return RequestError for expected client/policy failures. + BeforeDelete func(r *http.Request, streamName string) error } type Server struct { @@ -62,8 +71,9 @@ type Segment struct { idx int mutex *sync.Mutex cond *sync.Cond - buffer *bytes.Buffer + buffer *segmentBuffer closed bool + done atomic.Bool // to shut down any pending publishers closeCh chan bool @@ -83,6 +93,29 @@ const maxSegmentsPerStream = 5 var FirstByteTimeout = errors.New("pending read timeout") +// RequestError represents an expected request or policy failure from a hook. +type RequestError struct { + StatusCode int + Message string +} + +func (e *RequestError) Error() string { + if e.Message != "" { + return e.Message + } + if e.StatusCode != 0 { + return http.StatusText(e.StatusCode) + } + return http.StatusText(http.StatusBadRequest) +} + +func (e *RequestError) httpStatus() int { + if e.StatusCode >= 400 && e.StatusCode < 500 { + return e.StatusCode + } + return http.StatusBadRequest +} + func applyDefaults(config *TrickleServerConfig) { if config.BasePath == "" { config.BasePath = "/" @@ -259,6 +292,12 @@ func (sm *Server) closeStream(streamName string) error { func (sm *Server) handleDelete(w http.ResponseWriter, r *http.Request) { streamName := r.PathValue("streamName") + if sm.config.BeforeDelete != nil { + if err := sm.config.BeforeDelete(r, streamName); err != nil { + writeHookError(w, err) + return + } + } if err := sm.closeStream(streamName); err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return @@ -288,13 +327,30 @@ func (sm *Server) closeSeq(w http.ResponseWriter, r *http.Request) { } func (sm *Server) handleCreate(w http.ResponseWriter, r *http.Request) { - stream := sm.getOrCreateStream(r.PathValue("streamName"), r.Header.Get("Expect-Content"), false) + streamName := r.PathValue("streamName") + mimeType := r.Header.Get("Expect-Content") + if sm.config.BeforeCreate != nil { + if err := sm.config.BeforeCreate(r, streamName); err != nil { + writeHookError(w, err) + return + } + } + stream := sm.getOrCreateStream(streamName, mimeType, false) if stream == nil { http.Error(w, "Stream not found", http.StatusNotFound) return } } +func writeHookError(w http.ResponseWriter, err error) { + var requestErr *RequestError + if errors.As(err, &requestErr) { + http.Error(w, requestErr.Error(), requestErr.httpStatus()) + return + } + http.Error(w, err.Error(), http.StatusInternalServerError) +} + func (sm *Server) handlePost(w http.ResponseWriter, r *http.Request) { stream := sm.getOrCreateStream(r.PathValue("streamName"), r.Header.Get("Content-Type"), false) if stream == nil { @@ -374,6 +430,19 @@ func (tr *timeoutReader) Close() error { func (s *Stream) handlePost(w http.ResponseWriter, r *http.Request, idx int) { segment, _ := s.getForWrite(idx) + if r.Header.Get("Lp-Trickle-Reset") != "" { + // Usually means the publisher had to restart for some reason. + // Close prior segments to unblock subscribers for any hanging writes + // but allow for preconnected segments (sometimes they come out-of-order) + s.mutex.Lock() + for _, seg := range s.segments { + if seg != nil && seg.idx < segment.idx { + seg.close() + } + } + s.mutex.Unlock() + } + // Wrap the request body with the custom timeoutReader so we can send // provisional headers (keepalives) until receiving the first byte reader := &timeoutReader{ @@ -394,7 +463,7 @@ func (s *Stream) handlePost(w http.ResponseWriter, r *http.Request, idx int) { if totalRead == 0 { startedAt = time.Now() s.mutex.Lock() - s.nextWrite = idx + 1 + s.nextWrite = segment.idx + 1 s.writeTime = startedAt s.mutex.Unlock() } @@ -418,9 +487,10 @@ func (s *Stream) handlePost(w http.ResponseWriter, r *http.Request, idx int) { s.mutex.Lock() isClosed := s.closed // increment seq anyway: avoids clients erroring out on next seq - s.nextWrite = idx + 1 + s.nextWrite = segment.idx + 1 s.writeTime = startedAt s.mutex.Unlock() + w.Header().Set("Lp-Trickle-Seq", strconv.Itoa(segment.idx)) if isClosed { w.Header().Set("Lp-Trickle-Closed", "terminated") } @@ -440,6 +510,7 @@ func (s *Stream) handlePost(w http.ResponseWriter, r *http.Request, idx int) { } // Mark segment as closed + w.Header().Set("Lp-Trickle-Seq", strconv.Itoa(segment.idx)) segment.close() slog.Info("POST completed", "stream", s.name, "idx", idx, "bytes", totalRead, "took", time.Since(startedAt)) } @@ -588,14 +659,9 @@ func (s *Stream) handleGet(w http.ResponseWriter, r *http.Request, idx int) { latestSeq := s.nextWrite s.mutex.RUnlock() w.Header().Set("Lp-Trickle-Seq", strconv.Itoa(segment.idx)) + w.Header().Set("Lp-Trickle-Latest", strconv.Itoa(latestSeq)) if closed { w.Header().Set("Lp-Trickle-Closed", "terminated") - } else { - // usually happens if a publisher cancels a pending segment right before closing the channel - // other times, the subscriber is slow and the segment falls out of the live window - // send over latest seq so slow clients can grab leading edge - w.Header().Set("Lp-Trickle-Latest", strconv.Itoa(latestSeq)) - w.WriteHeader(470) } } return totalWrites, nil @@ -614,7 +680,7 @@ func newSegment(idx int) *Segment { mu := &sync.Mutex{} return &Segment{ idx: idx, - buffer: new(bytes.Buffer), + buffer: newSegmentBuffer(), cond: sync.NewCond(mu), mutex: mu, closeCh: make(chan bool), @@ -626,29 +692,43 @@ func (segment *Segment) writeData(data []byte) { defer segment.mutex.Unlock() // Write to buffer - segment.buffer.Write(data) + segment.buffer.write(data) // Signal waiting readers segment.cond.Broadcast() } -func (s *Segment) readData(startPos int) ([]byte, bool) { +func (s *Segment) readData(readPos int) ([]byte, int, bool) { + data, nextPos, atTail, invalid, _ := s.buffer.readChunk(readPos) + if invalid { + slog.Info("Invalid start pos, invoking eof") + return nil, readPos, true + } + if len(data) > 0 { + // A concurrent write+close can land after readChunk() snapshots published. + // Only signal EOF from the fast path if a fresh published load still + // matches the returned cursor. + if atTail && s.done.Load() && nextPos == int(s.buffer.published.Load()) { + return data, nextPos, true + } + return data, nextPos, false + } + s.mutex.Lock() defer s.mutex.Unlock() for { - totalLen := s.buffer.Len() - if startPos < totalLen { - data := s.buffer.Bytes()[startPos:totalLen] - return data, s.closed + data, nextPos, atTail, invalid, published := s.buffer.readChunk(readPos) + if len(data) > 0 { + return data, nextPos, s.closed && atTail && nextPos == published } - if startPos > totalLen { + if invalid { slog.Info("Invalid start pos, invoking eof") // This might happen if the buffer was reset // eg because of a repeated POST - return nil, true + return nil, readPos, true } if s.closed { - return nil, true + return nil, readPos, true } // Wait for new data s.cond.Wait() @@ -663,6 +743,7 @@ func (s *Segment) close() { defer s.mutex.Unlock() if !s.closed { s.closed = true + s.done.Store(true) close(s.closeCh) s.cond.Broadcast() } @@ -672,11 +753,11 @@ func (s *Segment) isFresh() bool { // fresh segments have not been written to yet s.mutex.Lock() defer s.mutex.Unlock() - return !s.closed && s.buffer.Len() == 0 + return !s.closed && s.buffer.isEmpty() } func (ss *SegmentSubscriber) readData() ([]byte, bool) { - data, eof := ss.segment.readData(ss.readPos) - ss.readPos += len(data) + data, nextPos, eof := ss.segment.readData(ss.readPos) + ss.readPos = nextPos return data, eof } diff --git a/trickle/trickle_test.go b/trickle/trickle_test.go index f5cd178904..bac63d9b9f 100644 --- a/trickle/trickle_test.go +++ b/trickle/trickle_test.go @@ -97,6 +97,319 @@ func TestTrickle_Close(t *testing.T) { require.Error(StreamNotFoundErr, pub2.Write(bytes.NewReader([]byte("bad post")))) } +func TestLocalPublisher_CreateContract(t *testing.T) { + t.Run("write without autocreate returns stream not found", func(t *testing.T) { + require := require.New(t) + server := ConfigureServer(TrickleServerConfig{ + Mux: http.NewServeMux(), // unused in practice with local-only publishing + Autocreate: false, + }) + pub := NewLocalPublisher(server, "missing", "text/plain") + + err := pub.Write(bytes.NewReader([]byte("hello"))) + + require.ErrorIs(err, StreamNotFoundErr) + _, exists := server.getStream("missing") + require.False(exists) + }) + + t.Run("create channel without autocreate then write succeeds", func(t *testing.T) { + require := require.New(t) + server := ConfigureServer(TrickleServerConfig{ + Mux: http.NewServeMux(), // unused in practice with local-only publishing + Autocreate: false, + }) + pub := NewLocalPublisher(server, "created", "text/plain") + + pub.CreateChannel() + err := pub.Write(bytes.NewReader([]byte("hello"))) + + require.NoError(err) + _, exists := server.getStream("created") + require.True(exists) + }) + + t.Run("write with autocreate creates missing channel", func(t *testing.T) { + require := require.New(t) + server := ConfigureServer(TrickleServerConfig{ + Mux: http.NewServeMux(), // unused in practice with local-only publishing + Autocreate: true, + }) + pub := NewLocalPublisher(server, "autocreated", "text/plain") + + err := pub.Write(bytes.NewReader([]byte("hello"))) + + require.NoError(err) + _, exists := server.getStream("autocreated") + require.True(exists) + }) +} + +func TestTrickle_HTTPCreateContract(t *testing.T) { + t.Run("create without autocreate returns not found", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: false, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/missing", "text/plain", nil) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusNotFound, resp.StatusCode) + }) + + t.Run("publish without autocreate returns not found", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: false, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/missing/0", "text/plain", bytes.NewReader([]byte("hello"))) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusNotFound, resp.StatusCode) + }) + + t.Run("create with autocreate succeeds", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/created", "text/plain", nil) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusOK, resp.StatusCode) + }) + + t.Run("publish with autocreate succeeds", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/published/0", "text/plain", bytes.NewReader([]byte("hello"))) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusOK, resp.StatusCode) + }) +} + +func TestTrickle_BeforeCreate(t *testing.T) { + t.Run("called before creation", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + var called bool + var server *Server + server = ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + BeforeCreate: func(r *http.Request, streamName string) error { + called = true + require.Equal("created", streamName) + require.Equal("text/plain", r.Header.Get("Expect-Content")) + _, exists := server.getStream(streamName) + require.False(exists) + return nil + }, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + req, err := http.NewRequest(http.MethodPost, ts.URL+"/created", nil) + require.NoError(err) + req.Header.Set("Expect-Content", "text/plain") + resp, err := http.DefaultClient.Do(req) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusOK, resp.StatusCode) + require.True(called) + resp, err = http.Get(ts.URL + "/created/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusOK, resp.StatusCode) + }) + + t.Run("request error prevents creation", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + BeforeCreate: func(r *http.Request, streamName string) error { + return &RequestError{StatusCode: http.StatusForbidden, Message: "nope"} + }, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/blocked", "text/plain", nil) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusForbidden, resp.StatusCode) + + resp, err = http.Get(ts.URL + "/blocked/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusNotFound, resp.StatusCode) + }) + + t.Run("generic error prevents creation", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + BeforeCreate: func(r *http.Request, streamName string) error { + return errors.New("boom") + }, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/errored", "text/plain", nil) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusInternalServerError, resp.StatusCode) + + resp, err = http.Get(ts.URL + "/errored/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusNotFound, resp.StatusCode) + }) +} + +func TestTrickle_BeforeDelete(t *testing.T) { + t.Run("called before deletion", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + var called bool + var server *Server + server = ConfigureServer(TrickleServerConfig{ + Mux: mux, + BeforeDelete: func(r *http.Request, streamName string) error { + called = true + require.Equal("deleted", streamName) + _, exists := server.getStream(streamName) + require.True(exists) + return nil + }, + }) + NewLocalPublisher(server, "deleted", "text/plain").CreateChannel() + ts := httptest.NewServer(mux) + defer ts.Close() + + req, err := http.NewRequest(http.MethodDelete, ts.URL+"/deleted", nil) + require.NoError(err) + resp, err := http.DefaultClient.Do(req) + require.NoError(err) + resp.Body.Close() + + require.Equal(http.StatusOK, resp.StatusCode) + require.True(called) + resp, err = http.Get(ts.URL + "/deleted/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusNotFound, resp.StatusCode) + }) + + t.Run("request error prevents deletion", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + server := ConfigureServer(TrickleServerConfig{ + Mux: mux, + BeforeDelete: func(r *http.Request, streamName string) error { + return &RequestError{StatusCode: http.StatusUnauthorized, Message: "nope"} + }, + }) + NewLocalPublisher(server, "blocked-delete", "text/plain").CreateChannel() + ts := httptest.NewServer(mux) + defer ts.Close() + + req, err := http.NewRequest(http.MethodDelete, ts.URL+"/blocked-delete", nil) + require.NoError(err) + resp, err := http.DefaultClient.Do(req) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusUnauthorized, resp.StatusCode) + + resp, err = http.Get(ts.URL + "/blocked-delete/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusOK, resp.StatusCode) + }) + + t.Run("generic error prevents deletion", func(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + server := ConfigureServer(TrickleServerConfig{ + Mux: mux, + BeforeDelete: func(r *http.Request, streamName string) error { + return errors.New("boom") + }, + }) + NewLocalPublisher(server, "errored-delete", "text/plain").CreateChannel() + ts := httptest.NewServer(mux) + defer ts.Close() + + req, err := http.NewRequest(http.MethodDelete, ts.URL+"/errored-delete", nil) + require.NoError(err) + resp, err := http.DefaultClient.Do(req) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusInternalServerError, resp.StatusCode) + + resp, err = http.Get(ts.URL + "/errored-delete/next") + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusOK, resp.StatusCode) + }) +} + +func TestTrickle_Delete(t *testing.T) { + require := require.New(t) + mux := http.NewServeMux() + ConfigureServer(TrickleServerConfig{ + Mux: mux, + Autocreate: true, + }) + ts := httptest.NewServer(mux) + defer ts.Close() + + resp, err := http.Post(ts.URL+"/delete-test", "text/plain", nil) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusOK, resp.StatusCode) + + req, err := http.NewRequest(http.MethodDelete, ts.URL+"/delete-test", nil) + require.NoError(err) + resp, err = http.DefaultClient.Do(req) + require.NoError(err) + resp.Body.Close() + require.Equal(http.StatusOK, resp.StatusCode) +} + func TestTrickle_SetSeq(t *testing.T) { require, channelURL := makeServer(t) @@ -152,7 +465,7 @@ func TestTrickle_Reset(t *testing.T) { require.Nil(err) wg := &sync.WaitGroup{} - // give preconnects time to latch on and autocreate the channel + // give preconnects time to latch on time.Sleep(5 * time.Millisecond) respCh := make(chan *http.Response) @@ -228,6 +541,113 @@ func TestTrickle_Reset(t *testing.T) { wg.Wait() } +// TestTrickle_PublisherReset verifies reset behavior for publisher restarts: +// 1) a blocked subscriber on an open segment is unblocked by reset, +// 2) already-written bytes on that segment remain readable, +// 3) reset write goes to the current nextWrite index. +func TestTrickle_PublisherReset(t *testing.T) { + require, channelURL, server := makeServerWithServer(t) + + lp := NewLocalPublisher(server, "testest", "text/plain") + lp.CreateChannel() + + // Partial write of segment 0 via pipe - do not close pipe yet. + r0, w0 := io.Pipe() + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _ = lp.Write(r0) + }() + _, err := w0.Write([]byte("Hello")) + require.Nil(err) + + // Local subscriber reads first bytes from segment 0, then blocks. + sub := NewLocalSubscriber(server, "testest") + sub.SetSeq(0) + td, err := sub.Read() + require.Nil(err) + + buf := make([]byte, 5) + _, err = io.ReadFull(td.Reader, buf) + require.Nil(err) + require.Equal("Hello", string(buf)) + + unblocked := make(chan []byte, 1) + go func() { + rest, _ := io.ReadAll(td.Reader) + unblocked <- rest + }() + + // HTTP POST reset - closes previous segments and writes next segment. + req, err := http.NewRequest("POST", channelURL+"/-1", bytes.NewReader([]byte("after-reset"))) + require.Nil(err) + req.Header.Set("Content-Type", "text/plain") + req.Header.Set("Lp-Trickle-Reset", "true") + resp, err := http.DefaultClient.Do(req) + require.Nil(err) + require.Equal(http.StatusOK, resp.StatusCode) + require.Equal("1", resp.Header.Get("Lp-Trickle-Seq"), "POST /-1 should echo resolved seq") + resp.Body.Close() + + // This receive deadlocks if reset doesn't unblock the reader. + <-unblocked + + // Re-read segment 0 - partial data should still be there. + sub.SetSeq(0) + td, err = sub.Read() + require.Nil(err) + data, err := io.ReadAll(td.Reader) + require.Nil(err) + require.Equal("Hello", string(data)) + + // Reset write should land at index 1. + td, err = sub.Read() + require.Nil(err) + data, err = io.ReadAll(td.Reader) + require.Nil(err) + require.Equal("after-reset", string(data)) + + // Additional writes to the old segment writer can still succeed until writer close. + // This is a bit racy w subscribers but subs will terminate once they catch up. + // NB: Fix this someday + _, err = w0.Write([]byte("late-bytes")) + require.NoError(err) + + require.Nil(w0.Close()) + <-writeDone +} + +func TestTrickle_EmptySegment(t *testing.T) { + require, channelURL := makeServer(t) + + pub, err := NewTricklePublisher(channelURL) + require.Nil(err) + defer pub.Close() + + pp, err := pub.Next() + require.Nil(err) + + n, err := pp.Write(bytes.NewReader(nil)) + require.Nil(err) + require.Equal(int64(0), n) + + sub, err := NewTrickleSubscriber(subConfig(t, channelURL)) + require.Nil(err) + sub.SetSeq(0) + + resp, err := sub.Read() + defer resp.Body.Close() + + require.Equal(http.StatusOK, resp.StatusCode) + require.Equal("0", resp.Header.Get("Lp-Trickle-Seq")) + require.Equal("", resp.Header.Get("Lp-Trickle-Closed")) + require.Equal("1", resp.Header.Get("Lp-Trickle-Latest")) + + body, err := io.ReadAll(resp.Body) + require.Nil(err) + require.Equal("", string(body)) +} + func TestTrickle_IdleSweep(t *testing.T) { require := require.New(t) mux := http.NewServeMux()