Exemple #1
0
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...)
}
Exemple #2
0
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...)
}
Exemple #3
0
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 := &registry.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 := &registry.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")
	}

}