dist: add methods and test

dist.go

@@ -303,6 +303,34 @@ func (d *Distributed) VNodeIndex(cell h3.H3Index) int {
 	return int(hashKey % d.vnodes)
 }
 
+// EachVNode iterate each vnode, calling fn for each vnode.
+func (d *Distributed) EachVNode(fn func(vnode uint64, addr string) bool) {
+	for i := uint64(0); i < d.vnodes; i++ {
+		addr, ok := d.Addr(i)
+		if !ok {
+			continue
+		}
+		if !fn(i, addr) {
+			break
+		}
+	}
+}
+
+// Addr returns the addr of the node by vnode id.
+func (d *Distributed) Addr(vnode uint64) (addr string, ok bool) {
+	hashKey := uint2hash(vnode)
+	idx := int(hashKey % d.vnodes)
+	d.mu.RLock()
+	node, found := d.index[idx]
+	d.mu.RUnlock()
+	if !found {
+		return
+	}
+	addr = node.addr
+	ok = true
+	return
+}
+
 // EachCell iterate each distributed cell, calling fn for each cell.
 func (d *Distributed) EachCell(iter func(c Cell)) {
 	d.mu.RLock()

dist_test.go

@@ -288,6 +288,66 @@ func TestDistributed_NeighborsFromLatLon(t *testing.T) {
 	}
 }
 
+func TestDistributed_EachVNode(t *testing.T) {
+	h3dist, err := New(Level3)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	hosts := []string{
+		"127.0.0.1",
+		"127.0.0.2",
+		"127.0.0.3",
+	}
+	for _, host := range hosts {
+		_ = h3dist.Add(host)
+	}
+	stats := make(map[string]int)
+
+	h3dist.EachVNode(func(id uint64, addr string) bool {
+		stats[addr]++
+		return true
+	})
+	for _, host := range hosts {
+		counter := stats[host]
+		if counter == 0 {
+			t.Fatalf("have %d, want > 0", counter)
+		}
+	}
+}
+
+func TestDistributed_Addr(t *testing.T) {
+	h3dist, err := New(Level3)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	hosts := []string{
+		"127.0.0.1",
+		"127.0.0.2",
+		"127.0.0.3",
+	}
+	for _, host := range hosts {
+		_ = h3dist.Add(host)
+	}
+
+	stats := make(map[string]int)
+	for i := uint64(0); i < h3dist.VNodes(); i++ {
+		addr, ok := h3dist.Addr(i)
+		if !ok {
+			continue
+		}
+		stats[addr]++
+	}
+
+	for _, host := range hosts {
+		counter := stats[host]
+		if counter == 0 {
+			t.Fatalf("have %d, want > 0", counter)
+		}
+	}
+}
+
 func TestDistributed_WhereIsMyParent(t *testing.T) {
 	h3dist, err := New(Level3)
 	if err != nil {