add Group.Upsert()

add Group.Upsert()

Author: werbenhu

group.go

@@ -70,6 +70,20 @@ func (b *Group) hashElement(element *Element) {
 	b.circle.Sort()
 }
 
+func (b *Group) Upsert(key string, payload []byte) error {
+	element := &Element{Key: key, Payload: payload}
+	b.Lock()
+	defer b.Unlock()
+	if _, ok := b.Elements[element.Key]; ok {
+		if err := b.delete(key); err != nil {
+			return err
+		}
+	}
+	b.Elements[element.Key] = element
+	b.hashElement(element)
+	return nil
+}
+
 func (b *Group) Insert(key string, payload []byte) error {
 	element := &Element{Key: key, Payload: payload}
 	b.Lock()
@@ -84,11 +98,8 @@ func (b *Group) Insert(key string, payload []byte) error {
 	return nil
 }
 
-func (b *Group) Delete(key string) error {
+func (b *Group) delete(key string) error {
 	element := &Element{Key: key, Payload: nil}
-	b.Lock()
-	defer b.Unlock()
-
 	delete(b.Elements, element.Key)
 	for i := 0; i < b.NumberOfReplicas; i++ {
 		virtualKey := b.virtualKey(key, i)
@@ -102,6 +113,12 @@ func (b *Group) Delete(key string) error {
 	return nil
 }
 
+func (b *Group) Delete(key string) error {
+	b.Lock()
+	defer b.Unlock()
+	return b.delete(key)
+}
+
 func (b *Group) Match(key string) (string, []byte, error) {
 	crc := b.hash(key)
 	b.RLock()

group_test.go

@@ -28,6 +28,27 @@ func TestGroupInit(t *testing.T) {
 	assert.NotNil(t, group.rows)
 }
 
+func TestGroupUpsert(t *testing.T) {
+	group := NewGroup("test", 10000)
+
+	key, payload := "192.168.1.100:1883", []byte("werbenhu100")
+	err := group.Upsert(key, payload)
+
+	assert.Nil(t, err)
+	assert.Equal(t, 10000, len(group.circle))
+	assert.Equal(t, 10000, len(group.rows))
+	assert.Equal(t, 1, len(group.Elements))
+
+	assert.Equal(t, key, group.Elements[key].Key)
+	assert.Equal(t, payload, group.Elements[key].Payload)
+
+	payload = []byte("werbenhu101")
+	err = group.Upsert(key, payload)
+	assert.Nil(t, err)
+	assert.Equal(t, key, group.Elements[key].Key)
+	assert.Equal(t, payload, group.Elements[key].Payload)
+}
+
 func TestGroupInsert(t *testing.T) {
 	group := NewGroup("test", 10000)
 
@@ -116,7 +137,7 @@ func TestGroupAll(t *testing.T) {
 }
 
 func TestGroupGetElements(t *testing.T) {
-	group := NewGroup("test", 10000)
+	group := NewGroup("testgetelements", 10000)
 	group.Insert("192.168.1.100:1883", []byte("werbenhu100"))
 	group.Insert("192.168.1.101:1883", []byte("werbenhu101"))