Merge pull request #839 from nbokovoy/pidset-lookup

Optimize PIDSet performance

actor/pidset.go

@@ -1,14 +1,18 @@
 package actor
 
-import "fmt"
-
 type PIDSet struct {
 	pids   []*PID
-	lookup map[string]*PID
+	lookup map[pidKey]int
+}
+
+// pidKey is used as a key in the lookup map to avoid allocations.
+type pidKey struct {
+	address string
+	id      string
 }
 
-func (p *PIDSet) key(pid *PID) string {
-	return fmt.Sprintf("%v:%v", pid.Address, pid.Id)
+func (p *PIDSet) key(pid *PID) pidKey {
+	return pidKey{address: pid.Address, id: pid.Id}
 }
 
 // NewPIDSet returns a new PIDSet with the given pids.
@@ -22,21 +26,21 @@ func NewPIDSet(pids ...*PID) *PIDSet {
 
 func (p *PIDSet) ensureInit() {
 	if p.lookup == nil {
-		p.lookup = make(map[string]*PID)
+		p.lookup = make(map[pidKey]int)
 	}
 }
 
 func (p *PIDSet) indexOf(v *PID) int {
-	for i, pid := range p.pids {
-		if v.Equal(pid) {
-			return i
-		}
+	if idx, ok := p.lookup[p.key(v)]; ok {
+		return idx
 	}
+
 	return -1
 }
 
 func (p *PIDSet) Contains(v *PID) bool {
-	return p.lookup[p.key(v)] != nil
+	_, ok := p.lookup[p.key(v)]
+	return ok
 }
 
 // Add adds the element v to the set.
@@ -45,8 +49,9 @@ func (p *PIDSet) Add(v *PID) {
 	if p.Contains(v) {
 		return
 	}
-	p.lookup[p.key(v)] = v
+
 	p.pids = append(p.pids, v)
+	p.lookup[p.key(v)] = len(p.pids) - 1
 }
 
 // Remove removes v from the set and returns true if them element existed.
@@ -58,8 +63,14 @@ func (p *PIDSet) Remove(v *PID) bool {
 	}
 
 	delete(p.lookup, p.key(v))
+	if i < len(p.pids)-1 {
+		lastPID := p.pids[len(p.pids)-1]
+
+		p.pids[i] = lastPID
+		p.lookup[p.key(lastPID)] = i
+	}
 
-	p.pids = append(p.pids[:i], p.pids[i+1:]...)
+	p.pids = p.pids[:len(p.pids)-1]
 
 	return true
 }
@@ -72,7 +83,7 @@ func (p *PIDSet) Len() int {
 // Clear removes all the elements in the set.
 func (p *PIDSet) Clear() {
 	p.pids = p.pids[:0]
-	p.lookup = make(map[string]*PID)
+	p.lookup = make(map[pidKey]int)
 }
 
 // Empty reports whether the set is empty.

actor/pidset_test.go

@@ -1,6 +1,7 @@
 package actor
 
 import (
+	"math/rand"
 	"strconv"
 	"testing"
 
@@ -23,6 +24,18 @@ func TestPIDSet_Clear(t *testing.T) {
 	assert.Len(t, s.pids, 0)
 }
 
+func TestPIDSet_Remove(t *testing.T) {
+	var s PIDSet
+	s.Add(NewPID(localAddress, "p1"))
+	s.Add(NewPID(localAddress, "p2"))
+	s.Add(NewPID(localAddress, "p3"))
+	assert.Equal(t, 3, s.Len())
+
+	s.Remove(NewPID(localAddress, "p3"))
+	assert.Equal(t, 2, s.Len())
+	assert.False(t, s.Contains(NewPID(localAddress, "p3")))
+}
+
 func TestPIDSet_AddSmall(t *testing.T) {
 	s := NewPIDSet()
 	p1 := NewPID(localAddress, "p1")
@@ -57,7 +70,7 @@ func TestPIDSet_AddMap(t *testing.T) {
 var pids []*PID
 
 func init() {
-	for i := 0; i < 1000; i++ {
+	for i := 0; i < 100000; i++ {
 		pids = append(pids, NewPID(localAddress, "p"+strconv.Itoa(i)))
 	}
 }
@@ -116,3 +129,36 @@ func pidSetAddRemove(b *testing.B, data []*PID) {
 		}
 	}
 }
+
+func BenchmarkPIDSet(b *testing.B) {
+	cases := []struct {
+		l int
+	}{
+		{l: 1},
+		{l: 5},
+		{l: 20},
+		{l: 500},
+		{l: 1000},
+		{l: 10000},
+		{l: 100000},
+	}
+
+	for _, tc := range cases {
+		b.Run("len "+strconv.Itoa(tc.l), func(b *testing.B) {
+			b.StopTimer()
+			var s PIDSet
+			for i := 0; i < tc.l; i++ {
+				s.Add(pids[i])
+			}
+			b.StartTimer()
+
+			for i := 0; i < b.N; i++ {
+				pid := pids[rand.Intn(len(pids))]
+
+				s.Add(pid)
+
+				s.Remove(s.Get(rand.Intn(s.Len())))
+			}
+		})
+	}
+}