Add some support packages

2 files changed, 380 insertions, 0 deletions

diff --git a/files/set.go b/files/set.go
new file mode 100644
index 0000000000..3d2af910f7
--- /dev/null
+++ b/files/set.go
@@ -0,0 +1,173 @@
+package fileset
+
+import "sync"
+
+type File struct {
+	Key      Key
+	Modified int64
+	Flags    uint32
+	Data     interface{}
+}
+
+type Key struct {
+	Name    string
+	Version uint32
+}
+
+type fileRecord struct {
+	Usage int
+	File  File
+}
+
+type bitset uint64
+
+func (a Key) newerThan(b Key) bool {
+	return a.Version > b.Version
+}
+
+type Set struct {
+	mutex              sync.RWMutex
+	files              map[Key]fileRecord
+	remoteKey          [64]map[string]Key
+	globalAvailability map[string]bitset
+	globalKey          map[string]Key
+}
+
+func NewSet() *Set {
+	var m = Set{
+		files:              make(map[Key]fileRecord),
+		globalAvailability: make(map[string]bitset),
+		globalKey:          make(map[string]Key),
+	}
+	return &m
+}
+
+func (m *Set) AddLocal(fs []File) {
+	m.mutex.Lock()
+	m.unlockedAddRemote(0, fs)
+	m.mutex.Unlock()
+}
+
+func (m *Set) SetLocal(fs []File) {
+	m.mutex.Lock()
+	m.unlockedSetRemote(0, fs)
+	m.mutex.Unlock()
+}
+
+func (m *Set) AddRemote(cid uint, fs []File) {
+	if cid < 1 || cid > 63 {
+		panic("Connection ID must be in the range 1 - 63 inclusive")
+	}
+	m.mutex.Lock()
+	m.unlockedAddRemote(cid, fs)
+	m.mutex.Unlock()
+}
+
+func (m *Set) SetRemote(cid uint, fs []File) {
+	if cid < 1 || cid > 63 {
+		panic("Connection ID must be in the range 1 - 63 inclusive")
+	}
+	m.mutex.Lock()
+	m.unlockedSetRemote(cid, fs)
+	m.mutex.Unlock()
+}
+
+func (m *Set) unlockedAddRemote(cid uint, fs []File) {
+	remFiles := m.remoteKey[cid]
+	for _, f := range fs {
+		n := f.Key.Name
+
+		if ck, ok := remFiles[n]; ok && ck == f.Key {
+			// The remote already has exactly this file, skip it
+			continue
+		}
+
+		remFiles[n] = f.Key
+
+		// Keep the block list or increment the usage
+		if br, ok := m.files[f.Key]; !ok {
+			m.files[f.Key] = fileRecord{
+				Usage: 1,
+				File:  f,
+			}
+		} else {
+			br.Usage++
+			m.files[f.Key] = br
+		}
+
+		// Update global view
+		gk, ok := m.globalKey[n]
+		switch {
+		case ok && f.Key == gk:
+			av := m.globalAvailability[n]
+			av |= 1 << cid
+			m.globalAvailability[n] = av
+		case f.Key.newerThan(gk):
+			m.globalKey[n] = f.Key
+			m.globalAvailability[n] = 1 << cid
+		}
+	}
+}
+
+func (m *Set) unlockedSetRemote(cid uint, fs []File) {
+	// Decrement usage for all files belonging to this remote, and remove
+	// those that are no longer needed.
+	for _, fk := range m.remoteKey[cid] {
+		br, ok := m.files[fk]
+		switch {
+		case ok && br.Usage == 1:
+			delete(m.files, fk)
+		case ok && br.Usage > 1:
+			br.Usage--
+			m.files[fk] = br
+		}
+	}
+
+	// Clear existing remote remoteKey
+	m.remoteKey[cid] = make(map[string]Key)
+
+	// Recalculate global based on all remaining remoteKey
+	for n := range m.globalKey {
+		var nk Key    // newest key
+		var na bitset // newest availability
+
+		for i, rem := range m.remoteKey {
+			if rk, ok := rem[n]; ok {
+				switch {
+				case rk == nk:
+					na |= 1 << uint(i)
+				case rk.newerThan(nk):
+					nk = rk
+					na = 1 << uint(i)
+				}
+			}
+		}
+
+		if na != 0 {
+			// Someone had the file
+			m.globalKey[n] = nk
+			m.globalAvailability[n] = na
+		} else {
+			// Noone had the file
+			delete(m.globalKey, n)
+			delete(m.globalAvailability, n)
+		}
+	}
+
+	// Add new remote remoteKey to the mix
+	m.unlockedAddRemote(cid, fs)
+}
+
+func (m *Set) Need(cid uint) []File {
+	var fs []File
+	m.mutex.Lock()
+
+	for name, gk := range m.globalKey {
+		if gk.newerThan(m.remoteKey[cid][name]) {
+			fs = append(fs, m.files[gk].File)
+		}
+	}
+
+	m.mutex.Unlock()
+	return fs
+}
diff --git a/files/set_test.go b/files/set_test.go
new file mode 100644
index 0000000000..0727df0f0f
--- /dev/null
+++ b/files/set_test.go
@@ -0,0 +1,207 @@
+package fileset
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+func TestGlobalSet(t *testing.T) {
+	m := NewSet()
+
+	local := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1000}, 0, 0, nil},
+		File{Key{"c", 1000}, 0, 0, nil},
+		File{Key{"d", 1000}, 0, 0, nil},
+	}
+
+	remote := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1001}, 0, 0, nil},
+		File{Key{"c", 1002}, 0, 0, nil},
+		File{Key{"e", 1000}, 0, 0, nil},
+	}
+
+	expectedGlobal := map[string]Key{
+		"a": local[0].Key,
+		"b": remote[1].Key,
+		"c": remote[2].Key,
+		"d": local[3].Key,
+		"e": remote[3].Key,
+	}
+
+	m.SetLocal(local)
+	m.SetRemote(1, remote)
+
+	if !reflect.DeepEqual(m.globalKey, expectedGlobal) {
+		t.Errorf("Global incorrect;\n%v !=\n%v", m.globalKey, expectedGlobal)
+	}
+
+	if lb := len(m.files); lb != 7 {
+		t.Errorf("Num files incorrect %d != 7\n%v", lb, m.files)
+	}
+}
+
+func BenchmarkSetLocal10k(b *testing.B) {
+	m := NewSet()
+
+	var local []File
+	for i := 0; i < 10000; i++ {
+		local = append(local, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	var remote []File
+	for i := 0; i < 10000; i++ {
+		remote = append(remote, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	m.SetRemote(1, remote)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		m.SetLocal(local)
+	}
+}
+
+func BenchmarkSetLocal10(b *testing.B) {
+	m := NewSet()
+
+	var local []File
+	for i := 0; i < 10; i++ {
+		local = append(local, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	var remote []File
+	for i := 0; i < 10000; i++ {
+		remote = append(remote, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	m.SetRemote(1, remote)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		m.SetLocal(local)
+	}
+}
+
+func BenchmarkAddLocal10k(b *testing.B) {
+	m := NewSet()
+
+	var local []File
+	for i := 0; i < 10000; i++ {
+		local = append(local, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	var remote []File
+	for i := 0; i < 10000; i++ {
+		remote = append(remote, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	m.SetRemote(1, remote)
+	m.SetLocal(local)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		for j := range local {
+			local[j].Key.Version++
+		}
+		b.StartTimer()
+		m.AddLocal(local)
+	}
+}
+
+func BenchmarkAddLocal10(b *testing.B) {
+	m := NewSet()
+
+	var local []File
+	for i := 0; i < 10; i++ {
+		local = append(local, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	var remote []File
+	for i := 0; i < 10000; i++ {
+		remote = append(remote, File{Key{fmt.Sprintf("file%d"), 1000}, 0, 0, nil})
+	}
+
+	m.SetRemote(1, remote)
+	m.SetLocal(local)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := range local {
+			local[j].Key.Version++
+		}
+		m.AddLocal(local)
+	}
+}
+
+func TestGlobalReset(t *testing.T) {
+	m := NewSet()
+
+	local := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1000}, 0, 0, nil},
+		File{Key{"c", 1000}, 0, 0, nil},
+		File{Key{"d", 1000}, 0, 0, nil},
+	}
+
+	remote := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1001}, 0, 0, nil},
+		File{Key{"c", 1002}, 0, 0, nil},
+		File{Key{"e", 1000}, 0, 0, nil},
+	}
+
+	expectedGlobalKey := map[string]Key{
+		"a": local[0].Key,
+		"b": local[1].Key,
+		"c": local[2].Key,
+		"d": local[3].Key,
+	}
+
+	m.SetLocal(local)
+	m.SetRemote(1, remote)
+	m.SetRemote(1, nil)
+
+	if !reflect.DeepEqual(m.globalKey, expectedGlobalKey) {
+		t.Errorf("Global incorrect;\n%v !=\n%v", m.globalKey, expectedGlobalKey)
+	}
+
+	if lb := len(m.files); lb != 4 {
+		t.Errorf("Num files incorrect %d != 4\n%v", lb, m.files)
+	}
+}
+
+func TestNeed(t *testing.T) {
+	m := NewSet()
+
+	local := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1000}, 0, 0, nil},
+		File{Key{"c", 1000}, 0, 0, nil},
+		File{Key{"d", 1000}, 0, 0, nil},
+	}
+
+	remote := []File{
+		File{Key{"a", 1000}, 0, 0, nil},
+		File{Key{"b", 1001}, 0, 0, nil},
+		File{Key{"c", 1002}, 0, 0, nil},
+		File{Key{"e", 1000}, 0, 0, nil},
+	}
+
+	shouldNeed := []File{
+		File{Key{"b", 1001}, 0, 0, nil},
+		File{Key{"c", 1002}, 0, 0, nil},
+		File{Key{"e", 1000}, 0, 0, nil},
+	}
+
+	m.SetLocal(local)
+	m.SetRemote(1, remote)
+
+	need := m.Need(0)
+	if !reflect.DeepEqual(need, shouldNeed) {
+		t.Errorf("Need incorrect;\n%v !=\n%v", need, shouldNeed)
+	}
+}