alice.pds.demo.boltless.dev
This repository has no description
1package mill
2
3import (
4 "context"
5 "errors"
6 "log/slog"
7 "sync"
8 "time"
9
10 millproto "tangled.org/core/spindle/mill/proto"
11 millv1 "tangled.org/core/spindle/mill/proto/gen"
12)
13
14var errSessionClosed = errors.New("mill: executor session closed")
15
16// one live websocket to an executor. many leases and async streams share
17// it, so one reader goroutine demuxes by message type and correlates by
18// lease ID. never hold locks across decodes
19type millSession struct {
20 nodeID string
21 epoch string
22 labels []string
23 enc messageEncoder
24 l *slog.Logger
25 closeTransport func() error
26
27 // snapshot, disconnected and graceTimer are guarded by Mill.mu, the
28 // fleet ranks across sessions under its own lock
29 snapshot *millv1.NodeSnapshot
30 disconnected bool
31 graceTimer *time.Timer
32 lastSeen time.Time
33
34 mu sync.Mutex
35 pending map[string]chan *millproto.Message // maps lease ID to response waiter
36
37 closeOnce sync.Once
38 closed chan struct{}
39}
40
41type messageEncoder interface {
42 Encode(*millproto.Message) error
43}
44
45func newSession(nodeID string, epoch string, labels []string, enc messageEncoder, l *slog.Logger) *millSession {
46 return &millSession{
47 nodeID: nodeID,
48 epoch: epoch,
49 labels: labels,
50 enc: enc,
51 l: l,
52 pending: make(map[string]chan *millproto.Message),
53 closed: make(chan struct{}),
54 lastSeen: time.Now(),
55 }
56}
57
58// caller holds Mill.mu
59func (s *millSession) live(grace time.Duration) bool {
60 return !s.disconnected && time.Since(s.lastSeen) <= grace
61}
62
63func (s *millSession) send(msg *millproto.Message) error {
64 return s.enc.Encode(msg)
65}
66
67func (s *millSession) close() {
68 s.closeOnce.Do(func() {
69 close(s.closed)
70 if s.closeTransport != nil {
71 _ = s.closeTransport()
72 }
73 })
74}
75
76// one-shot waiter for the next response on the lease, cancel unregisters it
77func (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) {
78 ch := make(chan *millproto.Message, 1)
79 s.mu.Lock()
80 s.pending[leaseID] = ch
81 s.mu.Unlock()
82 return ch, func() {
83 s.mu.Lock()
84 if s.pending[leaseID] == ch {
85 delete(s.pending, leaseID)
86 }
87 s.mu.Unlock()
88 }
89}
90
91func (s *millSession) deliver(leaseID string, msg *millproto.Message) {
92 s.mu.Lock()
93 ch := s.pending[leaseID]
94 delete(s.pending, leaseID)
95 s.mu.Unlock()
96 if ch != nil {
97 select {
98 case ch <- msg:
99 default:
100 }
101 }
102}
103
104// sends a message and waits for its response, respecting ctx and session closure
105func (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) {
106 if err := ctx.Err(); err != nil {
107 return nil, err
108 }
109 ch, cancel := s.await(leaseID)
110 defer cancel()
111
112 if err := s.send(msg); err != nil {
113 return nil, err
114 }
115
116 select {
117 case resp := <-ch:
118 return resp, nil
119 case <-ctx.Done():
120 return nil, ctx.Err()
121 case <-s.closed:
122 return nil, errSessionClosed
123 }
124}
125
126// demuxes frames until the decoder errors (connection gone)
127func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) error {
128 for {
129 msg, err := dec.Decode()
130 if err != nil {
131 return err
132 }
133 if err := s.dispatch(m, msg); err != nil {
134 return err
135 }
136 }
137}
138
139func (s *millSession) dispatch(m *Mill, msg *millproto.Message) error {
140 if !m.touchSession(s) {
141 return errSessionClosed
142 }
143 switch {
144 case msg.GetNodeSnapshot() != nil:
145 return m.onSnapshot(s, msg.GetNodeSnapshot())
146 case msg.GetReserveResult() != nil:
147 s.deliver(msg.GetReserveResult().GetLeaseId(), msg)
148 case msg.GetCommitted() != nil:
149 s.deliver(msg.GetCommitted().GetLeaseId(), msg)
150 case msg.GetEventBatch() != nil:
151 return m.onEventBatch(s, msg.GetEventBatch())
152 case msg.GetCancelAck() != nil:
153 m.onCancelAck(s, msg.GetCancelAck())
154 case msg.GetLiveLog() != nil:
155 return m.onLiveLog(s, msg.GetLiveLog())
156 default:
157 s.l.Warn("session received unexpected message", "node", s.nodeID)
158 }
159 return nil
160}