Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172package mill
import ( "context" "errors" "log/slog" "sync" "time"
millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen")
var errSessionClosed = errors.New("mill: executor session closed")
// one live websocket to an executor. many leases and async streams share// it, so one reader goroutine demuxes by message type and correlates by// lease ID. never hold locks across decodestype millSession struct { nodeID string epoch string arch string labels []string cacheStoreID string cacheNamespace string legacyProtocol bool enc messageEncoder l *slog.Logger closeTransport func() error
// recovery and cancellation state is guarded by Mill.mu snapshot *millv1.NodeSnapshot disconnected bool recovering bool recoveryReady bool recoveryExpired bool graceDeadline time.Time graceTimer *time.Timer lastSeen time.Time pendingCancels map[string]struct{}
mu sync.Mutex pending map[string]chan *millproto.Message // maps lease ID to response waiter
closeOnce sync.Once closed chan struct{}}
type messageEncoder interface { Encode(*millproto.Message) error}
func newSession(nodeID string, epoch string, labels []string, enc messageEncoder, l *slog.Logger) *millSession { return &millSession{ nodeID: nodeID, epoch: epoch, labels: labels, enc: enc, l: l, pending: make(map[string]chan *millproto.Message), pendingCancels: make(map[string]struct{}), closed: make(chan struct{}), lastSeen: time.Now(), }}
// caller holds Mill.mufunc (s *millSession) live(grace time.Duration) bool { return !s.disconnected && time.Since(s.lastSeen) <= grace}
func (s *millSession) send(msg *millproto.Message) error { return s.enc.Encode(msg)}
func (s *millSession) close() { s.closeOnce.Do(func() { close(s.closed) if s.closeTransport != nil { _ = s.closeTransport() } })}
// one-shot waiter for the next response on the lease, cancel unregisters itfunc (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) { ch := make(chan *millproto.Message, 1) s.mu.Lock() s.pending[leaseID] = ch s.mu.Unlock() return ch, func() { s.mu.Lock() if s.pending[leaseID] == ch { delete(s.pending, leaseID) } s.mu.Unlock() }}
func (s *millSession) deliver(leaseID string, msg *millproto.Message) { s.mu.Lock() ch := s.pending[leaseID] delete(s.pending, leaseID) s.mu.Unlock() if ch != nil { select { case ch <- msg: default: } }}
// sends a message and waits for its response, respecting ctx and session closurefunc (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) { if err := ctx.Err(); err != nil { return nil, err } ch, cancel := s.await(leaseID) defer cancel()
if err := s.send(msg); err != nil { return nil, err }
select { case resp := <-ch: return resp, nil case <-ctx.Done(): return nil, ctx.Err() case <-s.closed: return nil, errSessionClosed }}
// demuxes frames until the decoder errors (connection gone)func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) error { for { msg, err := dec.Decode() if err != nil { return err } if err := s.dispatch(m, msg); err != nil { return err } }}
func (s *millSession) dispatch(m *Mill, msg *millproto.Message) error { if !m.touchSession(s) { return errSessionClosed } switch { case msg.GetNodeSnapshot() != nil: return m.onSnapshot(s, msg.GetNodeSnapshot()) case msg.GetReserveResult() != nil: s.deliver(msg.GetReserveResult().GetLeaseId(), msg) case msg.GetCommitted() != nil: s.deliver(msg.GetCommitted().GetLeaseId(), msg) case msg.GetEventBatch() != nil: return m.onEventBatch(s, msg.GetEventBatch()) case msg.GetCancelAck() != nil: m.onCancelAck(s, msg.GetCancelAck()) case msg.GetLiveLog() != nil: return m.onLiveLog(s, msg.GetLiveLog()) case msg.GetQuotaReq() != nil: go m.handleQuotaRequest(s, msg.GetQuotaReq()) default: s.l.Warn("session received unexpected message", "node", s.nodeID) } return nil}