fix private host discovery (#18)

* fix host

* fix host lookup

* remove log

* fix check

* add another log

* more logs

* remove rtt check

* improve duplicate host cleanup

* remove logs

Author: luke-lombardi

pkg/client.go

@@ -349,6 +349,12 @@ func (c *BlobCacheClient) getGRPCClient(ctx context.Context, request *ClientRequ
 
 					if resp.Ok {
 						Logger.Infof("Content repopulated from source: %s\n", entry.SourcePath)
+						c.mu.Lock()
+						c.localHostCache[request.hash] = &localClientCache{
+							host:      host,
+							timestamp: time.Now(),
+						}
+						c.mu.Unlock()
 						return closestClient, host, nil
 					}
 

pkg/discovery.go

@@ -129,7 +129,7 @@ func (d *DiscoveryClient) discoverHostsViaTailscale(ctx context.Context) ([]*Blo
 }
 
 func (d *DiscoveryClient) discoverHostsViaMetadata(ctx context.Context) ([]*BlobCacheHost, error) {
-	hosts, err := d.metadata.GetAvailableHosts(ctx)
+	hosts, err := d.metadata.GetAvailableHosts(ctx, d.hostMap.Remove)
 	if err != nil {
 		return nil, err
 	}
@@ -140,18 +140,16 @@ func (d *DiscoveryClient) discoverHostsViaMetadata(ctx context.Context) ([]*Blob
 
 	for _, host := range hosts {
 		if host.PrivateAddr != "" {
-			addr := fmt.Sprintf("%s:%d", host.PrivateAddr, d.cfg.Port)
-
 			// Don't try to get the state on peers we're already aware of
-			if d.hostMap.Get(addr) != nil {
+			if d.hostMap.Get(host.Addr) != nil {
 				continue
 			}
 
 			wg.Add(1)
 			go func(addr string) {
 				defer wg.Done()
 
-				hostState, err := d.GetHostState(ctx, addr)
+				hostState, err := d.GetHostStateViaMetadata(ctx, addr, host.PrivateAddr)
 				if err != nil {
 					return
 				}
@@ -161,7 +159,7 @@ func (d *DiscoveryClient) discoverHostsViaMetadata(ctx context.Context) ([]*Blob
 				mu.Unlock()
 
 				Logger.Debugf("Added host with private address to map: %s", hostState.PrivateAddr)
-			}(addr)
+			}(host.Addr)
 		}
 	}
 
@@ -243,3 +241,37 @@ func (d *DiscoveryClient) GetHostState(ctx context.Context, addr string) (*BlobC
 
 	return &host, nil
 }
+
+func (d *DiscoveryClient) GetHostStateViaMetadata(ctx context.Context, addr, privateAddr string) (*BlobCacheHost, error) {
+	host := BlobCacheHost{
+		Addr:             addr,
+		RTT:              0,
+		PrivateAddr:      privateAddr,
+		CapacityUsagePct: 0.0,
+	}
+
+	var dialOpts = []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	}
+
+	conn, err := grpc.Dial(privateAddr, dialOpts...)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close()
+
+	c := proto.NewBlobCacheClient(conn)
+	resp, err := c.GetState(ctx, &proto.GetStateRequest{})
+	if err != nil {
+		return nil, err
+	}
+
+	host.RTT = 0
+	host.CapacityUsagePct = float64(resp.GetCapacityUsagePct())
+
+	if resp.GetVersion() != BlobCacheVersion {
+		return nil, fmt.Errorf("version mismatch: %s != %s", resp.GetVersion(), BlobCacheVersion)
+	}
+
+	return &host, nil
+}

pkg/metadata.go

@@ -287,7 +287,7 @@ func (m *BlobCacheMetadata) GetFsNodeChildren(ctx context.Context, id string) ([
 	return entries, nil
 }
 
-func (m *BlobCacheMetadata) GetAvailableHosts(ctx context.Context) ([]*BlobCacheHost, error) {
+func (m *BlobCacheMetadata) GetAvailableHosts(ctx context.Context, removeHostCallback func(host *BlobCacheHost)) ([]*BlobCacheHost, error) {
 	hostAddrs, err := m.rdb.SMembers(ctx, MetadataKeys.MetadataHostIndex()).Result()
 	if err != nil {
 		return nil, err
@@ -301,6 +301,7 @@ func (m *BlobCacheMetadata) GetAvailableHosts(ctx context.Context) ([]*BlobCache
 			// If the keepalive key doesn't exist, remove the host index key
 			if err == redis.Nil {
 				m.RemoveHostFromIndex(ctx, &BlobCacheHost{Addr: addr})
+				removeHostCallback(&BlobCacheHost{Addr: addr})
 			}
 
 			continue

pkg/server.go

@@ -60,7 +60,7 @@ func NewCacheService(ctx context.Context, cfg BlobCacheConfig) (*CacheService, e
 	if privateIpAddr != "" {
 		Logger.Infof("Discovered private ip address: %s", privateIpAddr)
 	}
-	currentHost.PrivateAddr = privateIpAddr
+	currentHost.PrivateAddr = fmt.Sprintf("%s:%d", privateIpAddr, cfg.Port)
 	currentHost.CapacityUsagePct = 0
 
 	cas, err := NewContentAddressableStorage(ctx, currentHost, metadata, cfg)
@@ -101,7 +101,7 @@ func NewCacheService(ctx context.Context, cfg BlobCacheConfig) (*CacheService, e
 func (cs *CacheService) HostKeepAlive() {
 	err := cs.metadata.AddHostToIndex(cs.ctx, cs.cas.currentHost)
 	if err != nil {
-		Logger.Fatalf("Failed to add host to index: %v", err)
+		Logger.Fatalf("Failed to add host to index? %v", err)
 	}
 
 	ticker := time.NewTicker(time.Duration(defaultHostKeepAliveIntervalS) * time.Second)
@@ -112,7 +112,7 @@ func (cs *CacheService) HostKeepAlive() {
 		case <-cs.ctx.Done():
 			return
 		case <-ticker.C:
-			cs.cas.currentHost.PrivateAddr = cs.privateIpAddr
+			cs.cas.currentHost.PrivateAddr = fmt.Sprintf("%s:%d", cs.privateIpAddr, cs.cfg.Port)
 			cs.cas.currentHost.CapacityUsagePct = cs.usagePct()
 			cs.metadata.SetHostKeepAlive(cs.ctx, cs.cas.currentHost)
 		}