github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/servermaster/scheduler/model/request.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package model
    15  
    16  import (
    17  	"github.com/pingcap/tiflow/engine/enginepb"
    18  	"github.com/pingcap/tiflow/engine/model"
    19  	resModel "github.com/pingcap/tiflow/engine/pkg/externalresource/model"
    20  	"github.com/pingcap/tiflow/pkg/errors"
    21  	"github.com/pingcap/tiflow/pkg/label"
    22  )
    23  
    24  // SchedulerRequest represents a request for an executor to run a given task.
    25  type SchedulerRequest struct {
    26  	TenantID          string // reserved for future use.
    27  	ExternalResources []resModel.ResourceKey
    28  	Selectors         []*label.Selector
    29  }
    30  
    31  // NewSchedulerRequestFromPB converts a protobuf message ScheduleTaskRequest
    32  // to the internal data structure SchedulerRequest.
    33  func NewSchedulerRequestFromPB(req *enginepb.ScheduleTaskRequest) (*SchedulerRequest, error) {
    34  	selectors := make([]*label.Selector, 0, len(req.GetSelectors()))
    35  	for _, sel := range req.GetSelectors() {
    36  		internalSel, err := SelectorFromPB(sel)
    37  		if err != nil {
    38  			return nil, err
    39  		}
    40  		selectors = append(selectors, internalSel)
    41  	}
    42  
    43  	schedulerReq := &SchedulerRequest{
    44  		ExternalResources: resModel.ToResourceKeys(req.GetResources()),
    45  		Selectors:         selectors,
    46  	}
    47  	return schedulerReq, nil
    48  }
    49  
    50  // SelectorFromPB converts a protobuf Selector message to the internal
    51  // data type label.Selector.
    52  // This function does the necessary sanity checks, so it's safe to use
    53  // with an RPC argument.
    54  func SelectorFromPB(req *enginepb.Selector) (*label.Selector, error) {
    55  	var op label.Op
    56  	switch req.Op {
    57  	case enginepb.Selector_Eq:
    58  		op = label.OpEq
    59  	case enginepb.Selector_Neq:
    60  		op = label.OpNeq
    61  	case enginepb.Selector_Regex:
    62  		op = label.OpRegex
    63  	default:
    64  		// Future-proof the code in case a new Op is added.
    65  		return nil, errors.ErrIncompatibleSchedulerRequest.GenWithStackByArgs("unknown selector op")
    66  	}
    67  
    68  	ret := &label.Selector{
    69  		Key:    label.Key(req.Label),
    70  		Target: req.Target,
    71  		Op:     op,
    72  	}
    73  
    74  	if err := ret.Validate(); err != nil {
    75  		return nil, errors.ErrIncompatibleSchedulerRequest.Wrap(err).GenWithStackByArgs("invalid selector")
    76  	}
    77  	return ret, nil
    78  }
    79  
    80  // SelectorToPB converts a label.Selector to a protobuf message Selector.
    81  func SelectorToPB(sel *label.Selector) (*enginepb.Selector, error) {
    82  	if err := sel.Validate(); err != nil {
    83  		return nil, errors.Annotate(err, "SelectorToPB")
    84  	}
    85  
    86  	var pbOp enginepb.Selector_Op
    87  	switch sel.Op {
    88  	case label.OpEq:
    89  		pbOp = enginepb.Selector_Eq
    90  	case label.OpNeq:
    91  		pbOp = enginepb.Selector_Neq
    92  	case label.OpRegex:
    93  		pbOp = enginepb.Selector_Regex
    94  	default:
    95  		return nil, errors.Errorf("unknown selector op %s", sel.Op)
    96  	}
    97  
    98  	return &enginepb.Selector{
    99  		Label:  string(sel.Key),
   100  		Target: sel.Target,
   101  		Op:     pbOp,
   102  	}, nil
   103  }
   104  
   105  // SchedulerResponse represents a response to a task scheduling request.
   106  type SchedulerResponse struct {
   107  	ExecutorID model.ExecutorID
   108  }