github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/plan/function/ctl/dist.go (about)

     1  // Copyright 2022 Matrix Origin
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ctl
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/clusterservice"
    19  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    20  	"github.com/matrixorigin/matrixone/pkg/pb/api"
    21  	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
    22  	"github.com/matrixorigin/matrixone/pkg/pb/txn"
    23  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    24  )
    25  
    26  // GetTNHandlerFunc used to handle dn's debug command handle func.
    27  // method: debug command type.
    28  // whichDN: used to decide which DNs to send the debug request to, nil returned means send all dns.
    29  // payload: used to get debug command request payload
    30  // repsonseUnmarshaler: used to unmarshal response
    31  func GetTNHandlerFunc(method api.OpCode,
    32  	whichTN func(parameter string) ([]uint64, error),
    33  	payload func(tnShardID uint64, parameter string, proc *process.Process) ([]byte, error),
    34  	repsonseUnmarshaler func([]byte) (any, error)) handleFunc {
    35  	return func(proc *process.Process,
    36  		service serviceType,
    37  		parameter string,
    38  		sender requestSender) (Result, error) {
    39  		if service != tn {
    40  			return Result{}, moerr.NewNotSupported(proc.Ctx, "service %s not supported", service)
    41  		}
    42  		targetTNs, err := whichTN(parameter)
    43  		if err != nil {
    44  			return Result{}, moerr.ConvertGoError(proc.Ctx, err)
    45  		}
    46  
    47  		containsTN := func(id uint64) bool {
    48  			for _, v := range targetTNs {
    49  				if v == id {
    50  					return true
    51  				}
    52  			}
    53  			return false
    54  		}
    55  
    56  		cluster := clusterservice.GetMOCluster()
    57  		var requests []txn.CNOpRequest
    58  		cluster.GetTNService(clusterservice.NewSelector(),
    59  			func(store metadata.TNService) bool {
    60  				for _, shard := range store.Shards {
    61  					if len(targetTNs) == 0 || containsTN(shard.ShardID) {
    62  						payLoad, e := payload(shard.ShardID, parameter, proc)
    63  						if e != nil {
    64  							err = e
    65  							return false
    66  						}
    67  						requests = append(requests, txn.CNOpRequest{
    68  							OpCode: uint32(method),
    69  							Target: metadata.TNShard{
    70  								TNShardRecord: metadata.TNShardRecord{
    71  									ShardID: shard.ShardID,
    72  								},
    73  								ReplicaID: shard.ReplicaID,
    74  								Address:   store.TxnServiceAddress,
    75  							},
    76  							Payload: payLoad,
    77  						})
    78  					}
    79  				}
    80  				return true
    81  			})
    82  		if err != nil {
    83  			return Result{}, err
    84  		}
    85  
    86  		results := make([]any, 0, len(requests))
    87  		if len(requests) > 0 {
    88  			responses, err := sender(proc.Ctx, proc, requests)
    89  			if err != nil {
    90  				return Result{}, err
    91  			}
    92  			if len(responses) != len(requests) {
    93  				panic("requests and response not match")
    94  			}
    95  
    96  			for _, resp := range responses {
    97  				r, err := repsonseUnmarshaler(resp.Payload)
    98  				if err != nil {
    99  					return Result{}, err
   100  				}
   101  				results = append(results, r)
   102  			}
   103  		}
   104  		// remove "Op" prefix
   105  		methodName, ok := api.OpMethodName[method]
   106  		if !ok {
   107  			return Result{Method: method.String(), Data: results}, nil
   108  		}
   109  		return Result{Method: methodName, Data: results}, nil
   110  	}
   111  }