Merge pull request #367 from jwasinger/pss-fwd-cache-cleaning

Add PSS fwd cache cleaning

swarm/pss/pss.go

@@ -177,7 +177,10 @@ func (self *Pss) Start(srv *p2p.Server) error {
 	go func() {
 		for {
 			tickC := time.Tick(defaultCleanInterval)
+			cacheTickC := time.Tick(self.cacheTTL)
 			select {
+			case <-cacheTickC:
+				self.cleanFwdCache()
 			case <-tickC:
 				self.cleanKeys()
 			case <-self.quitC:
@@ -756,6 +759,17 @@ func (self *Pss) forward(msg *PssMsg) {
 // SECTION: Caching
 /////////////////////////////////////////////////////////////////////
 
+// cleanFwdCache is used to periodically remove expired entries from the forward cache
+func (self *Pss) cleanFwdCache() {
+	self.fwdCacheMu.Lock()
+	defer self.fwdCacheMu.Unlock()
+	for k, v := range self.fwdCache {
+		if v.expiresAt.Before(time.Now()) {
+			delete(self.fwdCache, k)
+		}
+	}
+}
+
 // add a message to the cache
 func (self *Pss) addFwdCache(msg *PssMsg) error {
 	var entry pssCacheEntry

swarm/pss/pss_test.go

@@ -147,6 +147,7 @@ func TestCache(t *testing.T) {
 	pp := NewPssParams(privkey)
 	data := []byte("foo")
 	datatwo := []byte("bar")
+	datathree := []byte("baz")
 	wparams := &whisper.MessageParams{
 		TTL:      defaultWhisperTTL,
 		Src:      privkey,
@@ -169,6 +170,13 @@ func TestCache(t *testing.T) {
 		Payload: envtwo,
 		To:      to,
 	}
+	wparams.Payload = datathree
+	woutmsg, err = whisper.NewSentMessage(wparams)
+	envthree, err := woutmsg.Wrap(wparams)
+	msgthree := &PssMsg{
+		Payload: envthree,
+		To:      to,
+	}
 
 	digest := ps.digest(msg)
 	if err != nil {
@@ -178,6 +186,10 @@ func TestCache(t *testing.T) {
 	if err != nil {
 		t.Fatalf("could not store cache msgtwo: %v", err)
 	}
+	digestthree := ps.digest(msgthree)
+	if err != nil {
+		t.Fatalf("could not store cache msgthree: %v", err)
+	}
 
 	if digest == digesttwo {
 		t.Fatalf("different msgs return same hash: %d", digesttwo)
@@ -197,10 +209,23 @@ func TestCache(t *testing.T) {
 		t.Fatalf("message %v should NOT have EXPIRE record in cache but checkCache returned true", msgtwo)
 	}
 
-	time.Sleep(pp.CacheTTL)
+	time.Sleep(pp.CacheTTL + 1*time.Second)
+	err = ps.addFwdCache(msgthree)
+	if err != nil {
+		t.Fatalf("write to pss expire cache failed: %v", err)
+	}
+
 	if ps.checkFwdCache(msg) {
 		t.Fatalf("message %v should have expired from cache but checkCache returned true", msg)
 	}
+
+	if _, ok := ps.fwdCache[digestthree]; !ok {
+		t.Fatalf("unexpired message should be in the cache: %v", digestthree)
+	}
+
+	if _, ok := ps.fwdCache[digesttwo]; ok {
+		t.Fatalf("expired message should have been cleared from the cache: %v", digesttwo)
+	}
 }
 
 // matching of address hints; whether a message could be or is for the node
@@ -1309,6 +1334,7 @@ func newTestPss(privkey *ecdsa.PrivateKey, overlay network.Overlay, ppextra *Pss
 		pp.SymKeyCacheCapacity = ppextra.SymKeyCacheCapacity
 	}
 	ps := NewPss(overlay, pp)
+	ps.Start(nil)
 
 	return ps
 }