Merge pull request #38 from libp2p/feat/upsert-tag

Implement UpsertTag

p2p/net/connmgr/connmgr.go

@@ -244,6 +244,23 @@ func (cm *BasicConnMgr) UntagPeer(p peer.ID, tag string) {
 	delete(pi.tags, tag)
 }
 
+// UpsertTag is called to insert/update a peer tag
+func (cm *BasicConnMgr) UpsertTag(p peer.ID, tag string, upsert func(int) int) {
+	cm.lk.Lock()
+	defer cm.lk.Unlock()
+
+	pi, ok := cm.peers[p]
+	if !ok {
+		log.Info("tried to upsert tag from untracked peer: ", p)
+		return
+	}
+
+	oldval := pi.tags[tag]
+	newval := upsert(oldval)
+	pi.value += (newval - oldval)
+	pi.tags[tag] = newval
+}
+
 // CMInfo holds the configuration for BasicConnMgr, as well as status data.
 type CMInfo struct {
 	// The low watermark, as described in NewConnManager.

p2p/net/connmgr/connmgr_test.go

@@ -521,3 +521,35 @@ func TestPeerProtectionIdempotent(t *testing.T) {
 		t.Error("expected no protections")
 	}
 }
+
+func TestUpsertTag(t *testing.T) {
+	cm := NewConnManager(1, 1, time.Duration(10*time.Minute))
+	not := cm.Notifee()
+	conn := randConn(t, nil)
+	not.Connected(nil, conn)
+	rp := conn.RemotePeer()
+
+	cm.UpsertTag(rp, "tag", func(v int) int { return v + 1 })
+	if len(cm.peers[rp].tags) != 1 {
+		t.Fatal("expected a tag")
+	}
+	if cm.peers[rp].value != 1 {
+		t.Fatal("expected a tag value of 1")
+	}
+
+	cm.UpsertTag(rp, "tag", func(v int) int { return v + 1 })
+	if len(cm.peers[rp].tags) != 1 {
+		t.Fatal("expected a tag")
+	}
+	if cm.peers[rp].value != 2 {
+		t.Fatal("expected a tag value of 2")
+	}
+
+	cm.UpsertTag(rp, "tag", func(v int) int { return v - 1 })
+	if len(cm.peers[rp].tags) != 1 {
+		t.Fatal("expected a tag")
+	}
+	if cm.peers[rp].value != 1 {
+		t.Fatal("expected a tag value of 1")
+	}
+}