func (l *labelWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
md, _ := metadata.FromContext(ctx)
labels := make(map[string]string)
for k, v := range md {
if strings.HasPrefix(k, LabelPrefix) {
key := strings.TrimPrefix(k, LabelPrefix)
key = strings.ToLower(key)
labels[key] = strings.ToLower(v)
}
}
filter := func(services []*registry.Service) []*registry.Service {
var filtered []*registry.Service
SERVICE:
for _, service := range services {
// check labels for the service
for lk, lv := range labels {
// if we match then append the service and return
if v := service.Metadata[lk]; v == lv {
filtered = append(filtered, service)
continue SERVICE
}
}
var nodes []*registry.Node
NODE:
for _, node := range service.Nodes {
// check labels against the node
for lk, lv := range labels {
// if it matches then append node
if v := node.Metadata[lk]; v == lv {
nodes = append(nodes, node)
continue NODE
}
}
}
if len(nodes) > 0 {
service.Nodes = nodes
filtered = append(filtered, service)
}
}
return filtered
}
callOptions := append(opts, client.WithSelectOption(
selector.WithFilter(filter),
))
return l.Client.Call(ctx, req, rsp, callOptions...)
}
func (dc *dcWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
md, _ := metadata.FromContext(ctx)
filter := func(services []*registry.Service) []*registry.Service {
for _, service := range services {
var nodes []*registry.Node
for _, node := range service.Nodes {
if node.Metadata["datacenter"] == md["datacenter"] {
nodes = append(nodes, node)
}
}
service.Nodes = nodes
}
return services
}
callOptions := append(opts, client.WithSelectOption(
selector.WithFilter(filter),
))
fmt.Printf("[DC Wrapper] filtering for datacenter %s\n", md["datacenter"])
return dc.Client.Call(ctx, req, rsp, callOptions...)
}
func TestWatcher(t *testing.T) {
r := setupRegistry()
c := cache.NewSelector(selector.Registry(r))
// wait for watcher to get setup
time.Sleep(time.Millisecond)
// check that service is blank
if _, err := c.Select("foo.service"); err != registry.ErrNotFound {
log.Fatal("expected registry.ErrNotFound")
}
// setup svc
svc1 := ®istry.Service{Name: "foo.service", Version: "1"}
register(r, "pod-1", svc1)
time.Sleep(time.Millisecond)
var wg sync.WaitGroup
wg.Add(3)
c.Select("foo.service", selector.WithFilter(func(svcs []*registry.Service) []*registry.Service {
defer wg.Done()
if !hasServices(svcs, []*registry.Service{svc1}) {
t.Fatal("expected services to match")
}
return nil
}))
// setup svc
svc2 := ®istry.Service{Name: "foo.service", Version: "1"}
register(r, "pod-2", svc2)
time.Sleep(time.Millisecond)
c.Select("foo.service", selector.WithFilter(func(svcs []*registry.Service) []*registry.Service {
defer wg.Done()
if !hasNodes(svcs[0].Nodes, []*registry.Node{svc1.Nodes[0], svc2.Nodes[0]}) {
t.Fatal("expected to have same nodes")
}
return nil
}))
// deregister
os.Setenv("HOSTNAME", "pod-1")
r.Deregister(svc1)
time.Sleep(time.Millisecond)
c.Select("foo.service", selector.WithFilter(func(svcs []*registry.Service) []*registry.Service {
defer wg.Done()
if !hasServices(svcs, []*registry.Service{svc2}) {
t.Fatal("expected services to match")
}
return nil
}))
// remove pods
teardownRegistry()
time.Sleep(time.Millisecond)
if _, err := c.Select("foo.service"); err != registry.ErrNotFound {
log.Fatal("expected registry.ErrNotFound")
}
out := make(chan bool)
go func() {
wg.Wait()
close(out)
}()
select {
case <-out:
return
case <-time.After(time.Second):
t.Fatal("expected c.Select() to be called 3 times")
}
}