Method FrontendLoop

pkg/scheduler/scheduler.go:222–296

FrontendLoop handles connection from frontend.

(frontend schedulerpb.SchedulerForFrontend_FrontendLoopServer)

// FrontendLoop handles connection from frontend.
func (s *Scheduler) FrontendLoop(frontend schedulerpb.SchedulerForFrontend_FrontendLoopServer) error {
	frontendAddress, frontendCtx, err := s.frontendConnected(frontend)
	if err != nil {
		return err
	}
	defer s.frontendDisconnected(frontendAddress)

	// Response to INIT. If scheduler is not running, we skip for-loop, send SHUTTING_DOWN and exit this method.
	if s.State() == services.Running {
		if err := frontend.Send(&schedulerpb.SchedulerToFrontend{Status: schedulerpb.OK}); err != nil {
			return err
		}
	}

	// We stop accepting new queries in Stopping state. By returning quickly, we disconnect frontends, which in turns
	// cancels all their queries.
	for s.State() == services.Running {
		msg, err := frontend.Recv()
		if err != nil {
			// No need to report this as error, it is expected when query-frontend performs SendClose() (as frontendSchedulerWorker does).
			if err == io.EOF {
				return nil
			}
			return err
		}

		if s.State() != services.Running {
			break // break out of the loop, and send SHUTTING_DOWN message.
		}

		var resp *schedulerpb.SchedulerToFrontend

		switch msg.GetType() {
		case schedulerpb.ENQUEUE:

			// If there is a logical plan in the request body, we will fragment it before enqueueing
			// otherwise, it will be a single request and is the root and can be enqueued directly
			byteLP, err := getPlanFromHTTPRequest(msg.HttpRequest)
			if err != nil {
				return err
			}
			if len(byteLP) != 0 {
				err = s.fragmentAndEnqueueRequest(frontendCtx, frontendAddress, msg, byteLP)
			} else {
				err = s.enqueueRequest(frontendCtx, frontendAddress, msg, plan_fragments.Fragment{FragmentID: 0, IsRoot: true})
			}

			switch err {
			case nil:
				resp = &schedulerpb.SchedulerToFrontend{Status: schedulerpb.OK}
			case queue.ErrTooManyRequests:
				resp = &schedulerpb.SchedulerToFrontend{Status: schedulerpb.TOO_MANY_REQUESTS_PER_TENANT}
			default:
				resp = &schedulerpb.SchedulerToFrontend{Status: schedulerpb.ERROR, Error: err.Error()}
			}

		case schedulerpb.CANCEL:
			s.cancelRequestAndRemoveFromPending(frontendAddress, msg.QueryID, 0, true)