Repository: numaproj/numaflow
Pull request #169 (2022-09-13 07:06:58)

feat: Shuffle implementation
hash-based shuffle implementation with fnv
<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/pkg/shuffle/shuffle.go b/pkg/shuffle/shuffle.go
new file mode 100644
index 00000000..1805f20f
--- /dev/null
+++ b/pkg/shuffle/shuffle.go
@@ -0,0 +1,44 @@
+package shuffle
+
+import (
+	"github.com/numaproj/numaflow/pkg/isb"
+	"hash"
+	"hash/fnv"
+)
+
+// Shuffle shuffles messages among ISB
+type Shuffle struct {
+	bufferIdentifiers []string
+	buffersCount      int
+	hash              hash.Hash64
+}
+
+// NewShuffle accepts list of buffer identifiers(unique identifier of isb)
+// and returns new shuffle instance
+func NewShuffle(bufferIdentifiers []string) *Shuffle {
+	return &Shuffle{
+		bufferIdentifiers: bufferIdentifiers,
+		buffersCount:      len(bufferIdentifiers),
+		hash:              fnv.New64(),
+	}
+}
+
+// ShuffleMessages accepts list of isb messages and returns the mapping of isb to messages
+func (s *Shuffle) ShuffleMessages(messages []*isb.Message) map[string][]*isb.Message {
+
+	// hash of the message key returns a unique hashValue
+	// mod of hashValue will decide which isb it will belong
+	hashMap := make(map[string][]*isb.Message)
+	for _, message := range messages {
+		hashValue := s.generateHash(message.Key)
+		hashValue = hashValue % uint64(s.buffersCount)
+		hashMap[s.bufferIdentifiers[hashValue]] = append(hashMap[s.bufferIdentifiers[hashValue]], message)
+	}
+	return hashMap
+}
+
+func (s *Shuffle) generateHash(key string) uint64 {
+	s.hash.Reset()
+	_, _ = s.hash.Write([]byte(key))
+	return s.hash.Sum64()
+}


Test patch:
diff --git a/pkg/shuffle/shuffle_test.go b/pkg/shuffle/shuffle_test.go
new file mode 100644
index 00000000..5722b0eb
--- /dev/null
+++ b/pkg/shuffle/shuffle_test.go
@@ -0,0 +1,79 @@
+package shuffle
+
+import (
+	"fmt"
+	"github.com/numaproj/numaflow/pkg/isb"
+	"github.com/numaproj/numaflow/pkg/isb/testutils"
+	"github.com/stretchr/testify/assert"
+	"testing"
+	"time"
+)
+
+func TestShuffle_ShuffleMessages(t *testing.T) {
+
+	// buffer id list of test 1
+	bufferIdListOne := []string{
+		"buffer-1",
+		"buffer-2",
+		"buffer-3",
+		"buffer-4",
+	}
+
+	// buffer id list for test 2
+	var bufferIdListTwo []string
+
+	isbCount := 100
+	for i := 0; i < isbCount; i++ {
+		bufferIdListTwo = append(bufferIdListTwo, fmt.Sprintf("buffer-%d", i+1))
+	}
+
+	// build test messages for test 1
+	messagesOne := testutils.BuildTestWriteMessages(10000, time.Now())
+	// set key for messagesOne
+	var testMessagesOne []*isb.Message
+	for index := 0; index < len(messagesOne); index++ {
+		messagesOne[index].Key = fmt.Sprintf("key_%d", index)
+		testMessagesOne = append(testMessagesOne, &messagesOne[index])
+	}
+
+	// build test messages for test 2
+	messagesTwo := testutils.BuildTestWriteMessages(10, time.Now())
+	// set key for messages
+	var testMessagesTwo []*isb.Message
+	for index := 0; index < len(messagesTwo); index++ {
+		messagesOne[index].Key = fmt.Sprintf("key_%d", index)
+		testMessagesOne = append(testMessagesOne, &messagesTwo[index])
+	}
+
+	tests := []struct {
+		name               string
+		buffersIdentifiers []string
+		messages           []*isb.Message
+	}{
+		{
+			name:               "MessageCountGreaterThanBufferCount",
+			buffersIdentifiers: bufferIdListOne,
+			messages:           testMessagesOne,
+		},
+		{
+			name:               "BufferCountGreaterThanMessageCount",
+			buffersIdentifiers: bufferIdListTwo,
+			messages:           testMessagesTwo,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			// create shuffle with buffer id list
+			shuffler := NewShuffle(test.buffersIdentifiers)
+
+			bufferIdMessageMap := shuffler.ShuffleMessages(test.messages)
+			sum := 0
+			for _, value := range bufferIdMessageMap {
+				sum += len(value)
+			}
+
+			assert.Equal(t, sum, len(test.messages))
+		})
+	}
+}