go.etcd.io/etcd@v3.3.27+incompatible/proxy/grpcproxy/register.go (about)

     1  // Copyright 2017 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  package grpcproxy
    16  
    17  import (
    18  	"encoding/json"
    19  	"os"
    20  
    21  	"github.com/coreos/etcd/clientv3"
    22  	"github.com/coreos/etcd/clientv3/concurrency"
    23  	"github.com/coreos/etcd/clientv3/naming"
    24  
    25  	"golang.org/x/time/rate"
    26  	gnaming "google.golang.org/grpc/naming"
    27  )
    28  
    29  // allow maximum 1 retry per second
    30  const registerRetryRate = 1
    31  
    32  // Register registers itself as a grpc-proxy server by writing prefixed-key
    33  // with session of specified TTL (in seconds). The returned channel is closed
    34  // when the client's context is canceled.
    35  func Register(c *clientv3.Client, prefix string, addr string, ttl int) <-chan struct{} {
    36  	rm := rate.NewLimiter(rate.Limit(registerRetryRate), registerRetryRate)
    37  
    38  	donec := make(chan struct{})
    39  	go func() {
    40  		defer close(donec)
    41  
    42  		for rm.Wait(c.Ctx()) == nil {
    43  			ss, err := registerSession(c, prefix, addr, ttl)
    44  			if err != nil {
    45  				plog.Warningf("failed to create a session %v", err)
    46  				continue
    47  			}
    48  			select {
    49  			case <-c.Ctx().Done():
    50  				ss.Close()
    51  				return
    52  
    53  			case <-ss.Done():
    54  				plog.Warning("session expired; possible network partition or server restart")
    55  				plog.Warning("creating a new session to rejoin")
    56  				continue
    57  			}
    58  		}
    59  	}()
    60  
    61  	return donec
    62  }
    63  
    64  func registerSession(c *clientv3.Client, prefix string, addr string, ttl int) (*concurrency.Session, error) {
    65  	ss, err := concurrency.NewSession(c, concurrency.WithTTL(ttl))
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	gr := &naming.GRPCResolver{Client: c}
    71  	if err = gr.Update(c.Ctx(), prefix, gnaming.Update{Op: gnaming.Add, Addr: addr, Metadata: getMeta()}, clientv3.WithLease(ss.Lease())); err != nil {
    72  		return nil, err
    73  	}
    74  
    75  	plog.Infof("registered %q with %d-second lease", addr, ttl)
    76  	return ss, nil
    77  }
    78  
    79  // meta represents metadata of proxy register.
    80  type meta struct {
    81  	Name string `json:"name"`
    82  }
    83  
    84  func getMeta() string {
    85  	hostname, _ := os.Hostname()
    86  	bts, _ := json.Marshal(meta{Name: hostname})
    87  	return string(bts)
    88  }
    89  
    90  func decodeMeta(s string) (meta, error) {
    91  	m := meta{}
    92  	err := json.Unmarshal([]byte(s), &m)
    93  	return m, err
    94  }