github.com/rbisecke/kafka-go@v0.4.27/balancer_test.go (about)

     1  package kafka
     2  
     3  import (
     4  	"fmt"
     5  	"hash"
     6  	"hash/crc32"
     7  	"testing"
     8  )
     9  
    10  func TestHashBalancer(t *testing.T) {
    11  	testCases := map[string]struct {
    12  		Key        []byte
    13  		Hasher     hash.Hash32
    14  		Partitions []int
    15  		Partition  int
    16  	}{
    17  		"nil": {
    18  			Key:        nil,
    19  			Partitions: []int{0, 1, 2},
    20  			Partition:  0,
    21  		},
    22  		"partition-0": {
    23  			Key:        []byte("blah"),
    24  			Partitions: []int{0, 1},
    25  			Partition:  0,
    26  		},
    27  		"partition-1": {
    28  			Key:        []byte("blah"),
    29  			Partitions: []int{0, 1, 2},
    30  			Partition:  1,
    31  		},
    32  		"partition-2": {
    33  			Key:        []byte("boop"),
    34  			Partitions: []int{0, 1, 2},
    35  			Partition:  2,
    36  		},
    37  		"custom hash": {
    38  			Key:        []byte("boop"),
    39  			Hasher:     crc32.NewIEEE(),
    40  			Partitions: []int{0, 1, 2},
    41  			Partition:  1,
    42  		},
    43  		// in a previous version, this test would select a different partition
    44  		// than sarama's hash partitioner.
    45  		"hash code with MSB set": {
    46  			Key:        []byte("20"),
    47  			Partitions: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
    48  			Partition:  1,
    49  		},
    50  	}
    51  
    52  	for label, test := range testCases {
    53  		t.Run(label, func(t *testing.T) {
    54  			msg := Message{Key: test.Key}
    55  			h := Hash{
    56  				Hasher: test.Hasher,
    57  			}
    58  			partition := h.Balance(msg, test.Partitions...)
    59  			if partition != test.Partition {
    60  				t.Errorf("expected %v; got %v", test.Partition, partition)
    61  			}
    62  		})
    63  	}
    64  }
    65  
    66  func TestCRC32Balancer(t *testing.T) {
    67  	// These tests are taken from the default "consistent_random" partitioner from
    68  	// https://github.com/edenhill/librdkafka/blob/master/tests/0048-partitioner.c
    69  	partitionCount := 17
    70  	var partitions []int
    71  	for i := 0; i < partitionCount; i++ {
    72  		partitions = append(partitions, i*i)
    73  	}
    74  
    75  	testCases := map[string]struct {
    76  		Key        []byte
    77  		Partitions []int
    78  		Partition  int
    79  	}{
    80  		"nil": {
    81  			Key:        nil,
    82  			Partitions: partitions,
    83  			Partition:  -1,
    84  		},
    85  		"empty": {
    86  			Key:        []byte{},
    87  			Partitions: partitions,
    88  			Partition:  -1,
    89  		},
    90  		"unaligned": {
    91  			Key:        []byte("23456"),
    92  			Partitions: partitions,
    93  			Partition:  partitions[0xb1b451d7%partitionCount],
    94  		},
    95  		"long key": {
    96  			Key:        []byte("this is another string with more length to it perhaps"),
    97  			Partitions: partitions,
    98  			Partition:  partitions[0xb0150df7%partitionCount],
    99  		},
   100  		"short key": {
   101  			Key:        []byte("hejsan"),
   102  			Partitions: partitions,
   103  			Partition:  partitions[0xd077037e%partitionCount],
   104  		},
   105  	}
   106  
   107  	t.Run("default", func(t *testing.T) {
   108  		for label, test := range testCases {
   109  			t.Run(label, func(t *testing.T) {
   110  				b := CRC32Balancer{}
   111  				b.random.mock = -1
   112  
   113  				msg := Message{Key: test.Key}
   114  				partition := b.Balance(msg, test.Partitions...)
   115  				if partition != test.Partition {
   116  					t.Errorf("expected %v; got %v", test.Partition, partition)
   117  				}
   118  			})
   119  		}
   120  	})
   121  
   122  	t.Run("consistent", func(t *testing.T) {
   123  		b := CRC32Balancer{Consistent: true}
   124  		b.random.mock = -1
   125  
   126  		p := b.Balance(Message{}, partitions...)
   127  		if p < 0 {
   128  			t.Fatal("should not have gotten a random partition")
   129  		}
   130  		for i := 0; i < 10; i++ {
   131  			if p != b.Balance(Message{}, partitions...) {
   132  				t.Fatal("nil key should always hash consistently")
   133  			}
   134  			if p != b.Balance(Message{Key: []byte{}}, partitions...) {
   135  				t.Fatal("empty key should always hash consistently and have same result as nil key")
   136  			}
   137  		}
   138  	})
   139  }
   140  
   141  func TestMurmur2(t *testing.T) {
   142  	// These tests are taken from the "murmur2" implementation from
   143  	// https://github.com/edenhill/librdkafka/blob/master/src/rdmurmur2.c
   144  	testCases := []struct {
   145  		Key               []byte
   146  		JavaMurmur2Result uint32
   147  	}{
   148  		{Key: []byte("kafka"), JavaMurmur2Result: 0xd067cf64},
   149  		{Key: []byte("giberish123456789"), JavaMurmur2Result: 0x8f552b0c},
   150  		{Key: []byte("1234"), JavaMurmur2Result: 0x9fc97b14},
   151  		{Key: []byte("234"), JavaMurmur2Result: 0xe7c009ca},
   152  		{Key: []byte("34"), JavaMurmur2Result: 0x873930da},
   153  		{Key: []byte("4"), JavaMurmur2Result: 0x5a4b5ca1},
   154  		{Key: []byte("PreAmbleWillBeRemoved,ThePrePartThatIs"), JavaMurmur2Result: 0x78424f1c},
   155  		{Key: []byte("reAmbleWillBeRemoved,ThePrePartThatIs"), JavaMurmur2Result: 0x4a62b377},
   156  		{Key: []byte("eAmbleWillBeRemoved,ThePrePartThatIs"), JavaMurmur2Result: 0xe0e4e09e},
   157  		{Key: []byte("AmbleWillBeRemoved,ThePrePartThatIs"), JavaMurmur2Result: 0x62b8b43f},
   158  		{Key: []byte(""), JavaMurmur2Result: 0x106e08d9},
   159  		{Key: nil, JavaMurmur2Result: 0x106e08d9},
   160  	}
   161  
   162  	for _, test := range testCases {
   163  		t.Run(fmt.Sprintf("key:%s", test.Key), func(t *testing.T) {
   164  			got := murmur2(test.Key)
   165  			if got != test.JavaMurmur2Result {
   166  				t.Errorf("expected %v; got %v", test.JavaMurmur2Result, got)
   167  			}
   168  		})
   169  	}
   170  }
   171  
   172  func TestMurmur2Balancer(t *testing.T) {
   173  	// These tests are taken from the "murmur2_random" partitioner from
   174  	// https://github.com/edenhill/librdkafka/blob/master/tests/0048-partitioner.c
   175  	partitionCount := 17
   176  	librdkafkaPartitions := make([]int, partitionCount)
   177  	for i := 0; i < partitionCount; i++ {
   178  		librdkafkaPartitions[i] = i * i
   179  	}
   180  
   181  	// These tests are taken from the Murmur2Partitioner Python class from
   182  	// https://github.com/dpkp/kafka-python/blob/master/test/test_partitioner.py
   183  	pythonPartitions := make([]int, 1000)
   184  	for i := 0; i < 1000; i++ {
   185  		pythonPartitions[i] = i
   186  	}
   187  
   188  	testCases := map[string]struct {
   189  		Key        []byte
   190  		Partitions []int
   191  		Partition  int
   192  	}{
   193  		"librdkafka-nil": {
   194  			Key:        nil,
   195  			Partitions: librdkafkaPartitions,
   196  			Partition:  123,
   197  		},
   198  		"librdkafka-empty": {
   199  			Key:        []byte{},
   200  			Partitions: librdkafkaPartitions,
   201  			Partition:  librdkafkaPartitions[0x106e08d9%partitionCount],
   202  		},
   203  		"librdkafka-unaligned": {
   204  			Key:        []byte("23456"),
   205  			Partitions: librdkafkaPartitions,
   206  			Partition:  librdkafkaPartitions[0x058d780f%partitionCount],
   207  		},
   208  		"librdkafka-long key": {
   209  			Key:        []byte("this is another string with more length to it perhaps"),
   210  			Partitions: librdkafkaPartitions,
   211  			Partition:  librdkafkaPartitions[0x4f7703da%partitionCount],
   212  		},
   213  		"librdkafka-short key": {
   214  			Key:        []byte("hejsan"),
   215  			Partitions: librdkafkaPartitions,
   216  			Partition:  librdkafkaPartitions[0x5ec19395%partitionCount],
   217  		},
   218  		"python-empty": {
   219  			Key:        []byte(""),
   220  			Partitions: pythonPartitions,
   221  			Partition:  681,
   222  		},
   223  		"python-a": {
   224  			Key:        []byte("a"),
   225  			Partitions: pythonPartitions,
   226  			Partition:  524,
   227  		},
   228  		"python-ab": {
   229  			Key:        []byte("ab"),
   230  			Partitions: pythonPartitions,
   231  			Partition:  434,
   232  		},
   233  		"python-abc": {
   234  			Key:        []byte("abc"),
   235  			Partitions: pythonPartitions,
   236  			Partition:  107,
   237  		},
   238  		"python-123456789": {
   239  			Key:        []byte("123456789"),
   240  			Partitions: pythonPartitions,
   241  			Partition:  566,
   242  		},
   243  		"python-\x00 ": {
   244  			Key:        []byte{0, 32},
   245  			Partitions: pythonPartitions,
   246  			Partition:  742,
   247  		},
   248  	}
   249  
   250  	t.Run("default", func(t *testing.T) {
   251  		for label, test := range testCases {
   252  			t.Run(label, func(t *testing.T) {
   253  				b := Murmur2Balancer{}
   254  				b.random.mock = 123
   255  
   256  				msg := Message{Key: test.Key}
   257  				partition := b.Balance(msg, test.Partitions...)
   258  				if partition != test.Partition {
   259  					t.Errorf("expected %v; got %v", test.Partition, partition)
   260  				}
   261  			})
   262  		}
   263  	})
   264  
   265  	t.Run("consistent", func(t *testing.T) {
   266  		b := Murmur2Balancer{Consistent: true}
   267  		b.random.mock = -1
   268  
   269  		p := b.Balance(Message{}, librdkafkaPartitions...)
   270  		if p < 0 {
   271  			t.Fatal("should not have gotten a random partition")
   272  		}
   273  		for i := 0; i < 10; i++ {
   274  			if p != b.Balance(Message{}, librdkafkaPartitions...) {
   275  				t.Fatal("nil key should always hash consistently")
   276  			}
   277  		}
   278  	})
   279  }
   280  
   281  func TestLeastBytes(t *testing.T) {
   282  	testCases := map[string]struct {
   283  		Keys       [][]byte
   284  		Partitions [][]int
   285  		Partition  int
   286  	}{
   287  		"single message": {
   288  			Keys: [][]byte{
   289  				[]byte("key"),
   290  			},
   291  			Partitions: [][]int{
   292  				{0, 1, 2},
   293  			},
   294  			Partition: 0,
   295  		},
   296  		"multiple messages, no partition change": {
   297  			Keys: [][]byte{
   298  				[]byte("a"),
   299  				[]byte("ab"),
   300  				[]byte("abc"),
   301  				[]byte("abcd"),
   302  			},
   303  			Partitions: [][]int{
   304  				{0, 1, 2},
   305  				{0, 1, 2},
   306  				{0, 1, 2},
   307  				{0, 1, 2},
   308  			},
   309  			Partition: 0,
   310  		},
   311  		"partition gained": {
   312  			Keys: [][]byte{
   313  				[]byte("hello world 1"),
   314  				[]byte("hello world 2"),
   315  				[]byte("hello world 3"),
   316  			},
   317  			Partitions: [][]int{
   318  				{0, 1},
   319  				{0, 1},
   320  				{0, 1, 2},
   321  			},
   322  			Partition: 0,
   323  		},
   324  		"partition lost": {
   325  			Keys: [][]byte{
   326  				[]byte("hello world 1"),
   327  				[]byte("hello world 2"),
   328  				[]byte("hello world 3"),
   329  			},
   330  			Partitions: [][]int{
   331  				{0, 1, 2},
   332  				{0, 1, 2},
   333  				{0, 1},
   334  			},
   335  			Partition: 0,
   336  		},
   337  	}
   338  
   339  	for label, test := range testCases {
   340  		t.Run(label, func(t *testing.T) {
   341  			lb := &LeastBytes{}
   342  
   343  			var partition int
   344  			for i, key := range test.Keys {
   345  				msg := Message{Key: key}
   346  				partition = lb.Balance(msg, test.Partitions[i]...)
   347  			}
   348  
   349  			if partition != test.Partition {
   350  				t.Errorf("expected %v; got %v", test.Partition, partition)
   351  			}
   352  		})
   353  	}
   354  }