github.com/matrixorigin/matrixone@v1.2.0/pkg/proxy/plugin.go (about)

     1  // Copyright 2021 - 2023 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 proxy
    16  
    17  import (
    18  	"context"
    19  	v2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2"
    20  	"time"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    23  	"github.com/matrixorigin/matrixone/pkg/common/morpc"
    24  	"github.com/matrixorigin/matrixone/pkg/logutil"
    25  	"github.com/matrixorigin/matrixone/pkg/pb/plugin"
    26  )
    27  
    28  // pluginRouter is a router implementation that uses external plugin to select CN server.
    29  type pluginRouter struct {
    30  	// Router is a delegated impl that is used when plugin flags a Bypass action
    31  	Router
    32  	// plugin is the plugin that is used to select CN server
    33  	plugin Plugin
    34  }
    35  
    36  func newPluginRouter(r Router, p Plugin) *pluginRouter {
    37  	return &pluginRouter{
    38  		Router: r,
    39  		plugin: p,
    40  	}
    41  }
    42  
    43  // Route implements Router.Route.
    44  func (r *pluginRouter) Route(
    45  	ctx context.Context, ci clientInfo, filter func(uuid string) bool,
    46  ) (*CNServer, error) {
    47  	re, err := r.plugin.RecommendCN(ctx, ci)
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  	if re.Updated {
    52  		// plugin signals that a state updated has happened, request a refresh if the delegated router is refreshable
    53  		if rr, ok := r.Router.(RefreshableRouter); ok {
    54  			rr.Refresh(false)
    55  		}
    56  	}
    57  	switch re.Action {
    58  	case plugin.Select:
    59  		if re.CN == nil {
    60  			return nil, moerr.NewInternalErrorNoCtx("no CN server selected")
    61  		}
    62  		hash, err := ci.labelInfo.getHash()
    63  		if err != nil {
    64  			return nil, err
    65  		}
    66  		v2.ProxyConnectSelectCounter.Inc()
    67  		return &CNServer{
    68  			reqLabel: ci.labelInfo,
    69  			cnLabel:  re.CN.Labels,
    70  			uuid:     re.CN.ServiceID,
    71  			addr:     re.CN.SQLAddress,
    72  			hash:     hash,
    73  		}, nil
    74  	case plugin.Reject:
    75  		v2.ProxyConnectRejectCounter.Inc()
    76  		return nil, withCode(moerr.NewInfoNoCtx(re.Message),
    77  			codeAuthFailed)
    78  	case plugin.Bypass:
    79  		return r.Router.Route(ctx, ci, filter)
    80  	default:
    81  		return nil, moerr.NewInternalErrorNoCtx("unknown recommended action %d", re.Action)
    82  	}
    83  }
    84  
    85  // Plugin is the interface of proxy plugin.
    86  type Plugin interface {
    87  	// RecommendCN returns the recommended CN server.
    88  	RecommendCN(ctx context.Context, client clientInfo) (*plugin.Recommendation, error)
    89  }
    90  
    91  type rpcPlugin struct {
    92  	client  morpc.RPCClient
    93  	backend string
    94  	timeout time.Duration
    95  }
    96  
    97  func newRPCPlugin(backend string, timeout time.Duration) (*rpcPlugin, error) {
    98  	codec := morpc.NewMessageCodec(func() morpc.Message {
    99  		return &plugin.Response{}
   100  	})
   101  	backendOpts := []morpc.BackendOption{
   102  		morpc.WithBackendConnectTimeout(timeout),
   103  		morpc.WithBackendHasPayloadResponse(),
   104  		morpc.WithBackendLogger(logutil.GetGlobalLogger().Named("plugin-backend")),
   105  	}
   106  	bf := morpc.NewGoettyBasedBackendFactory(codec, backendOpts...)
   107  
   108  	clientOpts := []morpc.ClientOption{
   109  		morpc.WithClientInitBackends([]string{backend}, []int{1}),
   110  		morpc.WithClientMaxBackendPerHost(10),
   111  		morpc.WithClientLogger(logutil.GetGlobalLogger()),
   112  	}
   113  	cli, err := morpc.NewClient("plugin-client", bf, clientOpts...)
   114  	if err != nil {
   115  		return nil, err
   116  	}
   117  	return &rpcPlugin{client: cli, backend: backend, timeout: timeout}, nil
   118  }
   119  
   120  func (p *rpcPlugin) RecommendCN(ctx context.Context, ci clientInfo) (*plugin.Recommendation, error) {
   121  
   122  	resp, err := p.request(ctx, &plugin.Request{ClientInfo: &plugin.ClientInfo{
   123  		Tenant:        string(ci.Tenant),
   124  		Username:      ci.username,
   125  		OriginIP:      ci.originIP.String(),
   126  		LabelSelector: ci.labelInfo.allLabels(),
   127  	}})
   128  	if err != nil {
   129  		return nil, err
   130  	}
   131  	return resp.Recommendation, nil
   132  }
   133  
   134  func (p *rpcPlugin) Close() error {
   135  	return p.client.Close()
   136  }
   137  
   138  func (p *rpcPlugin) request(ctx context.Context, req *plugin.Request) (*plugin.Response, error) {
   139  	cc, cancel := context.WithTimeout(ctx, p.timeout)
   140  	defer cancel()
   141  	f, err := p.client.Send(cc, p.backend, req)
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  	defer f.Close()
   146  	resp, err := f.Get()
   147  	if err != nil {
   148  		return nil, err
   149  	}
   150  	return resp.(*plugin.Response), nil
   151  }