github.com/glide-im/glide@v1.6.0/pkg/rpc/selector.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"github.com/glide-im/glide/pkg/logger"
     6  	"github.com/smallnest/rpcx/client"
     7  	"github.com/smallnest/rpcx/share"
     8  )
     9  
    10  type selector struct {
    11  	services map[string]string
    12  	round    client.Selector
    13  	tags     map[string]string
    14  }
    15  
    16  func newSelector() *selector {
    17  	s := map[string]string{}
    18  	return &selector{
    19  		services: s,
    20  		round:    NewRoundRobinSelector(),
    21  		tags:     map[string]string{},
    22  	}
    23  }
    24  
    25  func (r *selector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
    26  
    27  	m := ctx.Value(share.ReqMetaDataKey).(map[string]string)
    28  
    29  	if target, ok := m["ExtraTarget"]; ok {
    30  		if _, ok := r.services[target]; ok {
    31  			return target
    32  		}
    33  		logger.E("unknown service addr, ExtraTarget:", target)
    34  	}
    35  
    36  	if tag, ok := m["ExtraTag"]; ok {
    37  		if path, ok := r.tags[tag]; ok {
    38  			if _, ok := r.services[path]; ok {
    39  				logger.D("route by tag: %s=%s", tag, path)
    40  				return path
    41  			}
    42  		}
    43  	}
    44  	return r.round.Select(ctx, servicePath, serviceMethod, args)
    45  }
    46  
    47  func (r *selector) UpdateServer(servers map[string]string) {
    48  	r.round.UpdateServer(servers)
    49  	for k, v := range servers {
    50  		r.services[k] = v
    51  	}
    52  }