github.com/lfch/etcd-io/tests/v3@v3.0.0-20221004140520-eac99acd3e9d/framework/integration/cluster_proxy.go (about)

     1  // Copyright 2016 The etcd Authors
     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  //go:build cluster_proxy
    16  // +build cluster_proxy
    17  
    18  package integration
    19  
    20  import (
    21  	"context"
    22  	"sync"
    23  
    24  	"github.com/lfch/etcd-io/client/v3"
    25  	"github.com/lfch/etcd-io/client/v3/namespace"
    26  	"github.com/lfch/etcd-io/server/v3/proxy/grpcproxy"
    27  	"github.com/lfch/etcd-io/server/v3/proxy/grpcproxy/adapter"
    28  )
    29  
    30  const ThroughProxy = true
    31  
    32  var (
    33  	pmu     sync.Mutex
    34  	proxies map[*clientv3.Client]grpcClientProxy = make(map[*clientv3.Client]grpcClientProxy)
    35  )
    36  
    37  const proxyNamespace = "proxy-namespace"
    38  
    39  type grpcClientProxy struct {
    40  	ctx       context.Context
    41  	ctxCancel func()
    42  	grpc      GrpcAPI
    43  	wdonec    <-chan struct{}
    44  	kvdonec   <-chan struct{}
    45  	lpdonec   <-chan struct{}
    46  }
    47  
    48  func ToGRPC(c *clientv3.Client) GrpcAPI {
    49  	pmu.Lock()
    50  	defer pmu.Unlock()
    51  
    52  	// dedicated context bound to 'grpc-proxy' lifetype
    53  	// (so in practice lifetime of the client connection to the proxy).
    54  	// TODO: Refactor to a separate clientv3.Client instance instead of the context alone.
    55  	ctx, ctxCancel := context.WithCancel(context.WithValue(context.TODO(), "_name", "grpcProxyContext"))
    56  
    57  	lg := c.GetLogger()
    58  
    59  	if v, ok := proxies[c]; ok {
    60  		return v.grpc
    61  	}
    62  
    63  	// test namespacing proxy
    64  	c.KV = namespace.NewKV(c.KV, proxyNamespace)
    65  	c.Watcher = namespace.NewWatcher(c.Watcher, proxyNamespace)
    66  	c.Lease = namespace.NewLease(c.Lease, proxyNamespace)
    67  	// test coalescing/caching proxy
    68  	kvp, kvpch := grpcproxy.NewKvProxy(c)
    69  	wp, wpch := grpcproxy.NewWatchProxy(ctx, lg, c)
    70  	lp, lpch := grpcproxy.NewLeaseProxy(ctx, c)
    71  	mp := grpcproxy.NewMaintenanceProxy(c)
    72  	clp, _ := grpcproxy.NewClusterProxy(lg, c, "", "") // without registering proxy URLs
    73  	authp := grpcproxy.NewAuthProxy(c)
    74  	lockp := grpcproxy.NewLockProxy(c)
    75  	electp := grpcproxy.NewElectionProxy(c)
    76  
    77  	grpc := GrpcAPI{
    78  		adapter.ClusterServerToClusterClient(clp),
    79  		adapter.KvServerToKvClient(kvp),
    80  		adapter.LeaseServerToLeaseClient(lp),
    81  		adapter.WatchServerToWatchClient(wp),
    82  		adapter.MaintenanceServerToMaintenanceClient(mp),
    83  		adapter.AuthServerToAuthClient(authp),
    84  		adapter.LockServerToLockClient(lockp),
    85  		adapter.ElectionServerToElectionClient(electp),
    86  	}
    87  	proxies[c] = grpcClientProxy{ctx: ctx, ctxCancel: ctxCancel, grpc: grpc, wdonec: wpch, kvdonec: kvpch, lpdonec: lpch}
    88  	return grpc
    89  }
    90  
    91  type proxyCloser struct {
    92  	clientv3.Watcher
    93  	proxyCtxCancel func()
    94  	wdonec         <-chan struct{}
    95  	kvdonec        <-chan struct{}
    96  	lclose         func()
    97  	lpdonec        <-chan struct{}
    98  }
    99  
   100  func (pc *proxyCloser) Close() error {
   101  	pc.proxyCtxCancel()
   102  	<-pc.kvdonec
   103  	err := pc.Watcher.Close()
   104  	<-pc.wdonec
   105  	pc.lclose()
   106  	<-pc.lpdonec
   107  	return err
   108  }
   109  
   110  func newClientV3(cfg clientv3.Config) (*clientv3.Client, error) {
   111  	c, err := clientv3.New(cfg)
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	rpc := ToGRPC(c)
   116  	c.KV = clientv3.NewKVFromKVClient(rpc.KV, c)
   117  	pmu.Lock()
   118  	lc := c.Lease
   119  	c.Lease = clientv3.NewLeaseFromLeaseClient(rpc.Lease, c, cfg.DialTimeout)
   120  	c.Watcher = &proxyCloser{
   121  		Watcher:        clientv3.NewWatchFromWatchClient(rpc.Watch, c),
   122  		wdonec:         proxies[c].wdonec,
   123  		kvdonec:        proxies[c].kvdonec,
   124  		lclose:         func() { lc.Close() },
   125  		lpdonec:        proxies[c].lpdonec,
   126  		proxyCtxCancel: proxies[c].ctxCancel,
   127  	}
   128  	pmu.Unlock()
   129  	return c, nil
   130  }