orchard

  1// Turns that outlive the tab that asked for them.
  2//
  3// A turn used to hang off the request context, so closing the tab, locking the
  4// phone or switching app cancelled it mid generation and lost everything it had
  5// already fetched. A turn takes up to twelve minutes and Isaac has one GPU, so
  6// the tab is the least reliable part of the arrangement and the wrong thing to
  7// tie the work to.
  8//
  9// So a turn runs detached and writes its events into a run, and a browser is
 10// only ever a reader. Closing the tab drops a reader. Opening the conversation
 11// again picks up a new one, which is handed everything the run has produced so
 12// far and then follows it live.
 13package main
 14
 15import (
 16	"context"
 17	"encoding/json"
 18	"log/slog"
 19	"sync"
 20	"time"
 21)
 22
 23// How long a finished run is kept so a tab that comes back can still read it.
 24// The messages are in the database by then, so this only has to cover the gap
 25// between finishing and the reader noticing.
 26const runKept = 30 * time.Minute
 27
 28// The ceiling on what one run holds. A turn's events are small and bounded by
 29// the tool rounds, and this only exists so a runaway cannot grow without limit.
 30const runMaxEvents = 4096
 31
 32type turnRun struct {
 33	mu sync.Mutex
 34	// Raw json rather than Event, because the last frame of a turn is the done
 35	// payload, which carries the stats, the sources and the rendered html and
 36	// is not an Event at all.
 37	events []json.RawMessage
 38	subs   map[chan json.RawMessage]struct{}
 39	done   bool
 40	// Set when the run ends, so a reader arriving after the fact is told the
 41	// turn is over rather than waiting on a stream that will never speak.
 42	endedAt time.Time
 43	cancel  context.CancelFunc
 44}
 45
 46// Runs holds every turn in flight, keyed by the conversation it belongs to. A
 47// new conversation has no id until its first turn is stored, so it is keyed by
 48// the id the browser was given up front instead.
 49type Runs struct {
 50	mu sync.Mutex
 51	m  map[string]*turnRun
 52}
 53
 54func NewRuns() *Runs { return &Runs{m: map[string]*turnRun{}} }
 55
 56// Start opens a run for a key, replacing and cancelling any run already there.
 57// Sending a second turn into the same conversation while one is still going
 58// means the first is no longer wanted, and leaving it running would have two
 59// generations writing into one conversation.
 60func (rs *Runs) Start(key string) *turnRun {
 61	rs.mu.Lock()
 62	defer rs.mu.Unlock()
 63	if old, ok := rs.m[key]; ok && old.cancel != nil {
 64		old.cancel()
 65	}
 66	r := &turnRun{subs: map[chan json.RawMessage]struct{}{}}
 67	rs.m[key] = r
 68	return r
 69}
 70
 71func (rs *Runs) Get(key string) (*turnRun, bool) {
 72	rs.mu.Lock()
 73	defer rs.mu.Unlock()
 74	r, ok := rs.m[key]
 75	return r, ok
 76}
 77
 78// Rekey moves a run to the conversation id the store gave it, so a tab that
 79// reopens the conversation by its real id finds the turn that is still running
 80// under the temporary one.
 81func (rs *Runs) Rekey(from, to string) {
 82	if from == to || to == "" {
 83		return
 84	}
 85	rs.mu.Lock()
 86	defer rs.mu.Unlock()
 87	if r, ok := rs.m[from]; ok {
 88		rs.m[to] = r
 89		delete(rs.m, from)
 90	}
 91}
 92
 93// Sweep drops runs that finished long enough ago that nobody is coming back for
 94// them. Called on a timer rather than on read, so a conversation nobody opens
 95// again does not sit in memory until the process restarts.
 96func (rs *Runs) Sweep(now time.Time) {
 97	rs.mu.Lock()
 98	defer rs.mu.Unlock()
 99	for k, r := range rs.m {
100		r.mu.Lock()
101		stale := r.done && now.Sub(r.endedAt) > runKept
102		r.mu.Unlock()
103		if stale {
104			delete(rs.m, k)
105		}
106	}
107}
108
109// Cancel stops a run, which is what a browser asking to stop a turn does. The
110// events already produced stay readable.
111func (rs *Runs) Cancel(key string) bool {
112	rs.mu.Lock()
113	r, ok := rs.m[key]
114	rs.mu.Unlock()
115	if !ok {
116		return false
117	}
118	r.mu.Lock()
119	cancel := r.cancel
120	r.mu.Unlock()
121	if cancel != nil {
122		cancel()
123	}
124	return true
125}
126
127func (r *turnRun) setCancel(cancel context.CancelFunc) {
128	r.mu.Lock()
129	defer r.mu.Unlock()
130	r.cancel = cancel
131}
132
133// Emit records an event and hands it to whoever is reading. A reader that is
134// not keeping up is skipped rather than waited on, since one slow tab must not
135// hold up the turn or the readers behind it.
136func (r *turnRun) Emit(v any) {
137	b, err := json.Marshal(v)
138	if err != nil {
139		slog.Error("an event would not marshal", "err", err)
140		return
141	}
142	r.mu.Lock()
143	defer r.mu.Unlock()
144	if r.done {
145		return
146	}
147	if len(r.events) < runMaxEvents {
148		r.events = append(r.events, b)
149	}
150	for ch := range r.subs {
151		select {
152		case ch <- b:
153		default:
154		}
155	}
156}
157
158// Finish closes the run to new events and wakes every reader so they can see it
159// ended rather than sitting on an open stream.
160func (r *turnRun) Finish() {
161	r.mu.Lock()
162	defer r.mu.Unlock()
163	if r.done {
164		return
165	}
166	r.done = true
167	r.endedAt = time.Now()
168	for ch := range r.subs {
169		close(ch)
170		delete(r.subs, ch)
171	}
172}
173
174// Follow hands back everything the run has already produced and a channel of
175// what comes next. The backlog is taken under the same lock that registers the
176// channel, so an event cannot land in the gap between the two and be lost.
177//
178// A closed channel means the turn is over. A run that had already finished
179// gives back its backlog and a closed channel, which is exactly what a tab
180// returning after the fact needs.
181func (r *turnRun) Follow() (backlog []json.RawMessage, ch <-chan json.RawMessage, live bool) {
182	r.mu.Lock()
183	defer r.mu.Unlock()
184	backlog = append([]json.RawMessage(nil), r.events...)
185	if r.done {
186		return backlog, nil, false
187	}
188	c := make(chan json.RawMessage, 256)
189	r.subs[c] = struct{}{}
190	return backlog, c, true
191}
192
193// Unfollow drops a reader, which happens when its tab goes away. The run does
194// not care how many readers it has and keeps going with none.
195func (r *turnRun) Unfollow(ch <-chan json.RawMessage) {
196	r.mu.Lock()
197	defer r.mu.Unlock()
198	for c := range r.subs {
199		if (<-chan json.RawMessage)(c) == ch {
200			delete(r.subs, c)
201			close(c)
202			return
203		}
204	}
205}
206
207// Running reports whether a turn is still going, which is what the conversation
208// endpoint tells a browser so it knows to attach rather than render and stop.
209func (r *turnRun) Running() bool {
210	r.mu.Lock()
211	defer r.mu.Unlock()
212	return !r.done
213}