| 296 | } |
| 297 | |
| 298 | func (s *Scheduler) frontendConnected(frontend schedulerpb.SchedulerForFrontend_FrontendLoopServer) (string, context.Context, error) { |
| 299 | msg, err := frontend.Recv() |
| 300 | if err != nil { |
| 301 | return "", nil, err |
| 302 | } |
| 303 | if msg.Type != schedulerpb.INIT || msg.FrontendAddress == "" { |
| 304 | return "", nil, errors.New("no frontend address") |
| 305 | } |
| 306 | |
| 307 | s.connectedFrontendsMu.Lock() |
| 308 | defer s.connectedFrontendsMu.Unlock() |
| 309 | |
| 310 | cf := s.connectedFrontends[msg.FrontendAddress] |
| 311 | if cf == nil { |
| 312 | cf = &connectedFrontend{ |
| 313 | connections: 0, |
| 314 | } |
| 315 | cf.ctx, cf.cancel = context.WithCancel(context.Background()) |
| 316 | s.connectedFrontends[msg.FrontendAddress] = cf |
| 317 | } |
| 318 | |
| 319 | cf.connections++ |
| 320 | return msg.FrontendAddress, cf.ctx, nil |
| 321 | } |
| 322 | |
| 323 | func (s *Scheduler) frontendDisconnected(frontendAddress string) { |
| 324 | s.connectedFrontendsMu.Lock() |