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") } }