github.com/matrixorigin/matrixone@v1.2.0/pkg/queryservice/client/query_client.go (about)

     1  // Copyright 2021 - 2024 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 client
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    21  	"github.com/matrixorigin/matrixone/pkg/common/morpc"
    22  	"github.com/matrixorigin/matrixone/pkg/common/runtime"
    23  	"github.com/matrixorigin/matrixone/pkg/defines"
    24  	pb "github.com/matrixorigin/matrixone/pkg/pb/query"
    25  )
    26  
    27  var methodVersions = map[pb.CmdMethod]int64{
    28  	pb.CmdMethod_ShowProcessList:       defines.MORPCVersion1,
    29  	pb.CmdMethod_AlterAccount:          defines.MORPCVersion1,
    30  	pb.CmdMethod_KillConn:              defines.MORPCVersion1,
    31  	pb.CmdMethod_TraceSpan:             defines.MORPCVersion1,
    32  	pb.CmdMethod_GetLockInfo:           defines.MORPCVersion1,
    33  	pb.CmdMethod_GetTxnInfo:            defines.MORPCVersion1,
    34  	pb.CmdMethod_GetCacheInfo:          defines.MORPCVersion1,
    35  	pb.CmdMethod_SyncCommit:            defines.MORPCVersion1,
    36  	pb.CmdMethod_GetCommit:             defines.MORPCVersion1,
    37  	pb.CmdMethod_RunTask:               defines.MORPCVersion1,
    38  	pb.CmdMethod_RemoveRemoteLockTable: defines.MORPCVersion1,
    39  	pb.CmdMethod_GetLatestBind:         defines.MORPCVersion1,
    40  	pb.CmdMethod_UnsubscribeTable:      defines.MORPCVersion1,
    41  	pb.CmdMethod_GetCacheData:          defines.MORPCVersion1,
    42  	pb.CmdMethod_GetStatsInfo:          defines.MORPCVersion1,
    43  	pb.CmdMethod_GetPipelineInfo:       defines.MORPCVersion1,
    44  	pb.CmdMethod_GetProtocolVersion:    defines.MORPCMinVersion, // To make sure these methods are compatible with all versions.
    45  	pb.CmdMethod_SetProtocolVersion:    defines.MORPCMinVersion,
    46  	pb.CmdMethod_CoreDumpConfig:        defines.MORPCMinVersion,
    47  	pb.CmdMethod_MigrateConnFrom:       defines.MORPCVersion1,
    48  	pb.CmdMethod_MigrateConnTo:         defines.MORPCVersion1,
    49  }
    50  
    51  type queryClient struct {
    52  	serviceID string
    53  	client    morpc.RPCClient
    54  	pool      morpc.MessagePool[*pb.Request, *pb.Response]
    55  }
    56  
    57  type QueryClient interface {
    58  	ServiceID() string
    59  	// SendMessage send message to a cache server.
    60  	SendMessage(ctx context.Context, address string, req *pb.Request) (*pb.Response, error)
    61  	// NewRequest creates a new request by cmd method.
    62  	NewRequest(pb.CmdMethod) *pb.Request
    63  	// Release releases the response.
    64  	Release(*pb.Response)
    65  	// Close closes the cache client.
    66  	Close() error
    67  }
    68  
    69  func NewQueryClient(sid string, cfg morpc.Config) (QueryClient, error) {
    70  	rt := runtime.ProcessLevelRuntime()
    71  	if rt == nil {
    72  		rt = runtime.DefaultRuntime()
    73  	}
    74  	pool := morpc.NewMessagePool(
    75  		func() *pb.Request { return &pb.Request{} },
    76  		func() *pb.Response { return &pb.Response{} },
    77  	)
    78  	c, err := cfg.NewClient("query-client", rt.Logger().Named("query-client").RawLogger(),
    79  		func() morpc.Message {
    80  			return pool.AcquireResponse()
    81  		},
    82  	)
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	cc := &queryClient{
    87  		serviceID: sid,
    88  		client:    c,
    89  		pool:      pool,
    90  	}
    91  	return cc, nil
    92  }
    93  
    94  func (c *queryClient) ServiceID() string {
    95  	return c.serviceID
    96  }
    97  
    98  // SendMessage implements the QueryService interface.
    99  func (c *queryClient) SendMessage(ctx context.Context, address string, req *pb.Request) (*pb.Response, error) {
   100  	if address == "" {
   101  		return nil, moerr.NewInternalError(ctx, "invalid CN query address %s", address)
   102  	}
   103  	if err := checkMethodVersion(ctx, req); err != nil {
   104  		return nil, err
   105  	}
   106  	f, err := c.client.Send(ctx, address, req)
   107  	if err != nil {
   108  		return nil, err
   109  	}
   110  	defer f.Close()
   111  	v, err := f.Get()
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	resp := v.(*pb.Response)
   116  	return c.unwrapResponseError(resp)
   117  }
   118  
   119  // NewRequest implements the QueryService interface.
   120  func (c *queryClient) NewRequest(method pb.CmdMethod) *pb.Request {
   121  	req := c.pool.AcquireRequest()
   122  	req.CmdMethod = method
   123  	return req
   124  }
   125  
   126  // Release implements the QueryService interface.
   127  func (c *queryClient) Release(resp *pb.Response) {
   128  	c.pool.ReleaseResponse(resp)
   129  }
   130  
   131  // Close implements the CacheService interface.
   132  func (c *queryClient) Close() error {
   133  	return c.client.Close()
   134  }
   135  
   136  func (c *queryClient) unwrapResponseError(resp *pb.Response) (*pb.Response, error) {
   137  	if err := resp.UnwrapError(); err != nil {
   138  		c.pool.ReleaseResponse(resp)
   139  		return nil, err
   140  	}
   141  	return resp, nil
   142  }
   143  
   144  func checkMethodVersion(ctx context.Context, req *pb.Request) error {
   145  	return runtime.CheckMethodVersion(ctx, methodVersions, req)
   146  }