From 0773ab3840809cafb6614d4e4aec1f4a12eb2faf Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Fri, 3 Jun 2016 13:52:28 -0700
Subject: [PATCH] fix cleanup of empty provider sets

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>
---
 routing/dht/providers.go      | 19 ++++++++++++++----
 routing/dht/providers_test.go | 38 +++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/routing/dht/providers.go b/routing/dht/providers.go
index a394123c8..09e263461 100644
--- a/routing/dht/providers.go
+++ b/routing/dht/providers.go
@@ -11,6 +11,9 @@ import (
 	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
 )
 
+var ProvideValidity = time.Hour * 24
+var defaultCleanupInterval = time.Hour
+
 type ProviderManager struct {
 	// all non channel fields are meant to be accessed only within
 	// the run method
@@ -23,6 +26,8 @@ type ProviderManager struct {
 	getprovs chan *getProv
 	period   time.Duration
 	proc     goprocess.Process
+
+	cleanupInterval time.Duration
 }
 
 type providerSet struct {
@@ -48,13 +53,14 @@ func NewProviderManager(ctx context.Context, local peer.ID) *ProviderManager {
 	pm.getlocal = make(chan chan []key.Key)
 	pm.local = make(map[key.Key]struct{})
 	pm.proc = goprocessctx.WithContext(ctx)
+	pm.cleanupInterval = defaultCleanupInterval
 	pm.proc.Go(func(p goprocess.Process) { pm.run() })
 
 	return pm
 }
 
 func (pm *ProviderManager) run() {
-	tick := time.NewTicker(time.Hour)
+	tick := time.NewTicker(pm.cleanupInterval)
 	for {
 		select {
 		case np := <-pm.newprovs:
@@ -85,16 +91,21 @@ func (pm *ProviderManager) run() {
 			lc <- keys
 
 		case <-tick.C:
-			for _, provs := range pm.providers {
+			for k, provs := range pm.providers {
 				var filtered []peer.ID
 				for p, t := range provs.set {
-					if time.Now().Sub(t) > time.Hour*24 {
+					if time.Now().Sub(t) > ProvideValidity {
 						delete(provs.set, p)
 					} else {
 						filtered = append(filtered, p)
 					}
 				}
-				provs.providers = filtered
+
+				if len(filtered) > 0 {
+					provs.providers = filtered
+				} else {
+					delete(pm.providers, k)
+				}
 			}
 
 		case <-pm.proc.Closing():
diff --git a/routing/dht/providers_test.go b/routing/dht/providers_test.go
index 8d25cfeba..9fa9d4b3a 100644
--- a/routing/dht/providers_test.go
+++ b/routing/dht/providers_test.go
@@ -2,6 +2,7 @@ package dht
 
 import (
 	"testing"
+	"time"
 
 	key "github.com/ipfs/go-ipfs/blocks/key"
 	peer "gx/ipfs/QmQGwpJy9P4yXZySmqkZEXCmbBpJUb8xntCv8Ca4taZwDC/go-libp2p-peer"
@@ -21,3 +22,40 @@ func TestProviderManager(t *testing.T) {
 	}
 	p.proc.Close()
 }
+
+func TestProvidesExpire(t *testing.T) {
+	ProvideValidity = time.Second
+	defaultCleanupInterval = time.Second
+
+	ctx := context.Background()
+	mid := peer.ID("testing")
+	p := NewProviderManager(ctx, mid)
+
+	peers := []peer.ID{"a", "b"}
+	var keys []key.Key
+	for i := 0; i < 10; i++ {
+		k := key.Key(i)
+		keys = append(keys, k)
+		p.AddProvider(ctx, k, peers[0])
+		p.AddProvider(ctx, k, peers[1])
+	}
+
+	for i := 0; i < 10; i++ {
+		out := p.GetProviders(ctx, keys[i])
+		if len(out) != 2 {
+			t.Fatal("expected providers to still be there")
+		}
+	}
+
+	time.Sleep(time.Second * 3)
+	for i := 0; i < 10; i++ {
+		out := p.GetProviders(ctx, keys[i])
+		if len(out) > 2 {
+			t.Fatal("expected providers to be cleaned up")
+		}
+	}
+
+	if len(p.providers) != 0 {
+		t.Fatal("providers map not cleaned up")
+	}
+}