diff --git a/server/ai_live_video.go b/server/ai_live_video.go index 5d38f44e17..63fc93a883 100644 --- a/server/ai_live_video.go +++ b/server/ai_live_video.go @@ -107,6 +107,11 @@ func startTricklePublish(ctx context.Context, url *url.URL, params aiRequestPara params.liveParams.segmentReader.SwitchReader(func(reader media.CloneableReader) { // check for end of stream if _, eos := reader.(*media.EOSReader); eos { + // Settle the final interval before publishing EOS so the gateway and + // orchestrator use the same accounting boundary. + if paymentProcessor != nil { + paymentProcessor.processSync(ctx) + } if err := publisher.Close(); err != nil { clog.Infof(ctx, "Error closing trickle publisher. err=%v", err) } diff --git a/server/live_payment_processor.go b/server/live_payment_processor.go index 1d7012bdfc..707c4ee923 100644 --- a/server/live_payment_processor.go +++ b/server/live_payment_processor.go @@ -21,6 +21,7 @@ type LivePaymentProcessor struct { lastProcessedAt time.Time lastProcessedMu sync.RWMutex + processMu sync.Mutex processCh chan time.Time processSegmentFunc func(units int64) error @@ -68,6 +69,9 @@ func (p *LivePaymentProcessor) start(ctx context.Context) { } func (p *LivePaymentProcessor) processOne(ctx context.Context, timestamp time.Time) { + p.processMu.Lock() + defer p.processMu.Unlock() + if p.shouldSkip(timestamp) { return } @@ -92,6 +96,10 @@ func (p *LivePaymentProcessor) processOne(ctx context.Context, timestamp time.Ti p.lastProcessedAt = p.lastProcessedAt.Add(processedDuration) } +func (p *LivePaymentProcessor) processSync(ctx context.Context) { + p.processOne(ctx, time.Now()) +} + func (p *LivePaymentProcessor) process(ctx context.Context) { timestamp := time.Now() if p.shouldSkip(timestamp) { diff --git a/server/live_payment_processor_test.go b/server/live_payment_processor_test.go index ea7c1fb14a..6e57678cd1 100644 --- a/server/live_payment_processor_test.go +++ b/server/live_payment_processor_test.go @@ -60,3 +60,21 @@ func TestLivePaymentProcessorCarriesFractionalUnits(t *testing.T) { require.Equal(t, int64(21), processed) }) } + +func TestLivePaymentProcessorProcessSyncCompletesBeforeCancel(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + var processed int64 + p := NewLivePaymentProcessor(ctx, time.Second, func(units int64) error { + processed += units + return nil + }) + + time.Sleep(time.Second) + p.processSync(ctx) + cancel() + synctest.Wait() + + require.Equal(t, int64(1), processed) + }) +}