|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "errors" |
| 5 | + "fmt" |
| 6 | + "net" |
| 7 | + "sync" |
| 8 | + "time" |
| 9 | + |
| 10 | + agent "github.com/livekit/protocol/livekit/agent" |
| 11 | + |
| 12 | + "github.com/livekit/livekit-cli/v2/pkg/ipc" |
| 13 | +) |
| 14 | + |
| 15 | +// devServer owns the dev-mode IPC channel between Go and the agent (Python) |
| 16 | +// processes. It listens on a loopback port that each (re)started agent connects |
| 17 | +// back to; every connection is handled by its own devSession. Reload (capturing |
| 18 | +// running jobs from the old process and restoring them in the new one) is one |
| 19 | +// feature served over this channel: |
| 20 | +// |
| 21 | +// 1. Go → old agent: GetRunningJobsRequest → GetRunningJobsResponse (capture) |
| 22 | +// 2. New agent → Go: GetRunningJobsRequest → Go replies with the saved jobs (restore) |
| 23 | +type devServer struct { |
| 24 | + listener *ipc.Listener |
| 25 | + mu sync.Mutex |
| 26 | + savedJobs *agent.GetRunningAgentJobsResponse |
| 27 | + |
| 28 | + // onServerInfo, if set, is invoked when a (re)started agent reports its |
| 29 | + // ServerInfo over the dev channel (agent name + the LiveKit URL it uses). |
| 30 | + onServerInfo func(agentName, url string) |
| 31 | +} |
| 32 | + |
| 33 | +func newDevServer() (*devServer, error) { |
| 34 | + ln, err := ipc.Listen("127.0.0.1:0") |
| 35 | + if err != nil { |
| 36 | + return nil, fmt.Errorf("dev server: %w", err) |
| 37 | + } |
| 38 | + return &devServer{listener: ln}, nil |
| 39 | +} |
| 40 | + |
| 41 | +func (rs *devServer) addr() string { |
| 42 | + return rs.listener.Addr().String() |
| 43 | +} |
| 44 | + |
| 45 | +// newSession wraps a freshly accepted dev-channel connection. The session |
| 46 | +// answers inbound GetRunningJobsRequests with the captured jobs (restore) and |
| 47 | +// forwards any ServerInfo notifications. |
| 48 | +func (rs *devServer) newSession(conn net.Conn) *devSession { |
| 49 | + s := newDevSession(conn) |
| 50 | + s.onServerInfo = rs.onServerInfo |
| 51 | + s.jobsProvider = rs.takeSavedJobs |
| 52 | + return s |
| 53 | +} |
| 54 | + |
| 55 | +// takeSavedJobs returns and clears the captured jobs. Jobs are restored to the |
| 56 | +// first process that asks; subsequent asks (within the same generation) get none. |
| 57 | +func (rs *devServer) takeSavedJobs() *agent.GetRunningAgentJobsResponse { |
| 58 | + rs.mu.Lock() |
| 59 | + defer rs.mu.Unlock() |
| 60 | + saved := rs.savedJobs |
| 61 | + rs.savedJobs = nil |
| 62 | + return saved |
| 63 | +} |
| 64 | + |
| 65 | +// captureJobs asks the (old) session for its running jobs and stores them so the |
| 66 | +// next process can restore them. Best-effort: failures are logged, not fatal. |
| 67 | +func (rs *devServer) captureJobs(s *devSession) { |
| 68 | + resp, err := s.getRunningJobs(1500 * time.Millisecond) |
| 69 | + if err != nil { |
| 70 | + out.Warnf("reload: failed to capture running jobs: %v", err) |
| 71 | + return |
| 72 | + } |
| 73 | + if resp != nil { |
| 74 | + rs.mu.Lock() |
| 75 | + rs.savedJobs = resp |
| 76 | + rs.mu.Unlock() |
| 77 | + out.Statusf("reload: captured %d running job(s)", len(resp.Jobs)) |
| 78 | + } |
| 79 | +} |
| 80 | + |
| 81 | +func (rs *devServer) close() error { |
| 82 | + return rs.listener.Close() |
| 83 | +} |
| 84 | + |
| 85 | +// devMsgKind identifies an outbound call awaiting a reply, keyed by the response |
| 86 | +// message type. Type-based routing is unambiguous today because the CLI never has |
| 87 | +// two outbound calls of the same response type in flight; add a correlation id to |
| 88 | +// AgentDevMessage if that ever changes. |
| 89 | +type devMsgKind int |
| 90 | + |
| 91 | +const ( |
| 92 | + kindJobsResponse devMsgKind = iota |
| 93 | +) |
| 94 | + |
| 95 | +var ( |
| 96 | + errDevSessionClosed = errors.New("dev session closed") |
| 97 | + errDevSessionTimeout = errors.New("dev session request timed out") |
| 98 | +) |
| 99 | + |
| 100 | +// devSession owns a single dev-channel IPC connection and multiplexes the |
| 101 | +// AgentDevMessage protocol over it. A single read loop (run) dispatches inbound |
| 102 | +// notifications (ServerInfo) and peer requests (GetRunningJobsRequest) to their |
| 103 | +// handlers, while outbound calls (getRunningJobs) await their matching response. |
| 104 | +// This keeps one owner on the connection so the CLI/agent API can grow new |
| 105 | +// request/response pairs and callbacks without per-message lifecycle juggling. |
| 106 | +type devSession struct { |
| 107 | + conn net.Conn |
| 108 | + |
| 109 | + // onServerInfo, if set, is invoked when the peer pushes a ServerInfo. |
| 110 | + onServerInfo func(agentName, url string) |
| 111 | + // jobsProvider, if set, supplies the response to an inbound |
| 112 | + // GetRunningJobsRequest (the restore side of a reload). |
| 113 | + jobsProvider func() *agent.GetRunningAgentJobsResponse |
| 114 | + |
| 115 | + writeMu sync.Mutex // serializes writes to conn |
| 116 | + |
| 117 | + mu sync.Mutex |
| 118 | + pending map[devMsgKind]chan *agent.AgentDevMessage |
| 119 | +} |
| 120 | + |
| 121 | +func newDevSession(conn net.Conn) *devSession { |
| 122 | + return &devSession{ |
| 123 | + conn: conn, |
| 124 | + pending: make(map[devMsgKind]chan *agent.AgentDevMessage), |
| 125 | + } |
| 126 | +} |
| 127 | + |
| 128 | +// run reads and dispatches messages until the connection closes. It is meant to |
| 129 | +// run in its own goroutine for the lifetime of the connection. |
| 130 | +func (s *devSession) run() { |
| 131 | + defer s.failPending() |
| 132 | + for { |
| 133 | + msg := &agent.AgentDevMessage{} |
| 134 | + if err := ipc.ReadProto(s.conn, msg); err != nil { |
| 135 | + return |
| 136 | + } |
| 137 | + switch msg.Message.(type) { |
| 138 | + case *agent.AgentDevMessage_ServerInfo: |
| 139 | + if s.onServerInfo != nil { |
| 140 | + si := msg.GetServerInfo() |
| 141 | + s.onServerInfo(si.GetAgentName(), si.GetUrl()) |
| 142 | + } |
| 143 | + case *agent.AgentDevMessage_GetRunningJobsRequest: |
| 144 | + s.handleJobsRequest() |
| 145 | + case *agent.AgentDevMessage_GetRunningJobsResponse: |
| 146 | + s.deliver(kindJobsResponse, msg) |
| 147 | + } |
| 148 | + // Unknown message types are ignored, leaving room for protocol growth. |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +// handleJobsRequest answers a peer's GetRunningJobsRequest with the jobs supplied |
| 153 | +// by jobsProvider (the restore side of a reload). |
| 154 | +func (s *devSession) handleJobsRequest() { |
| 155 | + var resp *agent.GetRunningAgentJobsResponse |
| 156 | + if s.jobsProvider != nil { |
| 157 | + resp = s.jobsProvider() |
| 158 | + } |
| 159 | + if resp == nil { |
| 160 | + resp = &agent.GetRunningAgentJobsResponse{} |
| 161 | + } |
| 162 | + err := s.write(&agent.AgentDevMessage{ |
| 163 | + Message: &agent.AgentDevMessage_GetRunningJobsResponse{GetRunningJobsResponse: resp}, |
| 164 | + }) |
| 165 | + if err != nil { |
| 166 | + out.Warnf("reload: failed to send restore response: %v", err) |
| 167 | + } else if len(resp.Jobs) > 0 { |
| 168 | + out.Statusf("reload: restored %d job(s) to new process", len(resp.Jobs)) |
| 169 | + } |
| 170 | +} |
| 171 | + |
| 172 | +// getRunningJobs asks the peer for its running jobs and waits up to timeout for |
| 173 | +// the reply (the capture side of a reload). |
| 174 | +func (s *devSession) getRunningJobs(timeout time.Duration) (*agent.GetRunningAgentJobsResponse, error) { |
| 175 | + ch := s.register(kindJobsResponse) |
| 176 | + defer s.unregister(kindJobsResponse) |
| 177 | + |
| 178 | + err := s.write(&agent.AgentDevMessage{ |
| 179 | + Message: &agent.AgentDevMessage_GetRunningJobsRequest{GetRunningJobsRequest: &agent.GetRunningAgentJobsRequest{}}, |
| 180 | + }) |
| 181 | + if err != nil { |
| 182 | + return nil, err |
| 183 | + } |
| 184 | + |
| 185 | + select { |
| 186 | + case msg, ok := <-ch: |
| 187 | + if !ok || msg == nil { |
| 188 | + return nil, errDevSessionClosed |
| 189 | + } |
| 190 | + return msg.GetGetRunningJobsResponse(), nil |
| 191 | + case <-time.After(timeout): |
| 192 | + return nil, errDevSessionTimeout |
| 193 | + } |
| 194 | +} |
| 195 | + |
| 196 | +func (s *devSession) write(msg *agent.AgentDevMessage) error { |
| 197 | + s.writeMu.Lock() |
| 198 | + defer s.writeMu.Unlock() |
| 199 | + return ipc.WriteProto(s.conn, msg) |
| 200 | +} |
| 201 | + |
| 202 | +func (s *devSession) register(kind devMsgKind) chan *agent.AgentDevMessage { |
| 203 | + ch := make(chan *agent.AgentDevMessage, 1) |
| 204 | + s.mu.Lock() |
| 205 | + s.pending[kind] = ch |
| 206 | + s.mu.Unlock() |
| 207 | + return ch |
| 208 | +} |
| 209 | + |
| 210 | +func (s *devSession) unregister(kind devMsgKind) { |
| 211 | + s.mu.Lock() |
| 212 | + delete(s.pending, kind) |
| 213 | + s.mu.Unlock() |
| 214 | +} |
| 215 | + |
| 216 | +// deliver routes an inbound response to the waiting caller, if any. |
| 217 | +func (s *devSession) deliver(kind devMsgKind, msg *agent.AgentDevMessage) { |
| 218 | + s.mu.Lock() |
| 219 | + ch := s.pending[kind] |
| 220 | + delete(s.pending, kind) |
| 221 | + s.mu.Unlock() |
| 222 | + if ch != nil { |
| 223 | + ch <- msg // buffered (cap 1): never blocks |
| 224 | + } |
| 225 | +} |
| 226 | + |
| 227 | +// failPending unblocks every outstanding caller when the connection closes. |
| 228 | +func (s *devSession) failPending() { |
| 229 | + s.mu.Lock() |
| 230 | + for kind, ch := range s.pending { |
| 231 | + close(ch) |
| 232 | + delete(s.pending, kind) |
| 233 | + } |
| 234 | + s.mu.Unlock() |
| 235 | +} |
| 236 | + |
| 237 | +func (s *devSession) close() error { |
| 238 | + return s.conn.Close() |
| 239 | +} |
0 commit comments