Send kafka notification messages in batches when queue_dir is enabled (#18164)

Fixes #18124

internal/store/batch.go

@@ -0,0 +1,117 @@
+// Copyright (c) 2023 MinIO, Inc.
+//
+// This file is part of MinIO Object Storage stack
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package store
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+)
+
+// ErrBatchFull indicates that the batch is full
+var ErrBatchFull = errors.New("batch is full")
+
+type key interface {
+	string | int | int64
+}
+
+// Batch represents an ordered batch
+type Batch[K key, T any] struct {
+	keys  []K
+	items map[K]T
+	limit uint32
+
+	sync.Mutex
+}
+
+// Add adds the item to the batch
+func (b *Batch[K, T]) Add(key K, item T) error {
+	b.Lock()
+	defer b.Unlock()
+
+	if b.isFull() {
+		return ErrBatchFull
+	}
+
+	if _, ok := b.items[key]; !ok {
+		b.keys = append(b.keys, key)
+	}
+	b.items[key] = item
+
+	return nil
+}
+
+// GetAll fetches the items and resets the batch
+// Returned items are not referenced by the batch
+func (b *Batch[K, T]) GetAll() (orderedKeys []K, orderedItems []T, err error) {
+	b.Lock()
+	defer b.Unlock()
+
+	orderedKeys = append([]K(nil), b.keys...)
+	for _, key := range orderedKeys {
+		item, ok := b.items[key]
+		if !ok {
+			err = fmt.Errorf("item not found for the key: %v; should not happen;", key)
+			return
+		}
+		orderedItems = append(orderedItems, item)
+		delete(b.items, key)
+	}
+
+	b.keys = b.keys[:0]
+
+	return
+}
+
+// GetByKey will get the batch item by the provided key
+func (b *Batch[K, T]) GetByKey(key K) (T, bool) {
+	b.Lock()
+	defer b.Unlock()
+
+	item, ok := b.items[key]
+	return item, ok
+}
+
+// Len returns the no of items in the batch
+func (b *Batch[K, T]) Len() int {
+	b.Lock()
+	defer b.Unlock()
+
+	return len(b.keys)
+}
+
+// IsFull checks if the batch is full or not
+func (b *Batch[K, T]) IsFull() bool {
+	b.Lock()
+	defer b.Unlock()
+
+	return b.isFull()
+}
+
+func (b *Batch[K, T]) isFull() bool {
+	return len(b.items) >= int(b.limit)
+}
+
+// NewBatch creates a new batch
+func NewBatch[K key, T any](limit uint32) *Batch[K, T] {
+	return &Batch[K, T]{
+		keys:  make([]K, 0, limit),
+		items: make(map[K]T, limit),
+		limit: limit,
+	}
+}

internal/store/batch_test.go

@@ -0,0 +1,126 @@
+// Copyright (c) 2023 MinIO, Inc.
+//
+// This file is part of MinIO Object Storage stack
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package store
+
+import (
+	"errors"
+	"sync"
+	"testing"
+)
+
+func TestBatch(t *testing.T) {
+	var limit uint32 = 100
+	batch := NewBatch[int, int](limit)
+	for i := 0; i < int(limit); i++ {
+		if err := batch.Add(i, i); err != nil {
+			t.Fatalf("failed to add %v; %v", i, err)
+		}
+		if _, ok := batch.GetByKey(i); !ok {
+			t.Fatalf("failed to get the item by key %v after adding", i)
+		}
+	}
+	err := batch.Add(101, 101)
+	if err == nil || !errors.Is(err, ErrBatchFull) {
+		t.Fatalf("Expected err %v but got %v", ErrBatchFull, err)
+	}
+	if !batch.IsFull() {
+		t.Fatal("Expected batch.IsFull to be true but got false")
+	}
+	batchLen := batch.Len()
+	if batchLen != int(limit) {
+		t.Fatalf("expected batch length to be %v but got %v", limit, batchLen)
+	}
+	keys, items, err := batch.GetAll()
+	if err != nil {
+		t.Fatalf("unable to get the items from the batch; %v", err)
+	}
+	if len(items) != int(limit) {
+		t.Fatalf("Expected length of the batch items to be %v but got %v", limit, len(items))
+	}
+	if len(keys) != int(limit) {
+		t.Fatalf("Expected length of the batch keys to be %v but got %v", limit, len(items))
+	}
+	batchLen = batch.Len()
+	if batchLen != 0 {
+		t.Fatalf("expected batch to be empty but still left with %d items", batchLen)
+	}
+	// Add duplicate entries
+	for i := 0; i < 10; i++ {
+		if err := batch.Add(99, 99); err != nil {
+			t.Fatalf("failed to add duplicate item %v to batch after Get; %v", i, err)
+		}
+	}
+	if _, ok := batch.GetByKey(99); !ok {
+		t.Fatal("failed to get the duplicxate item by key '99' after adding")
+	}
+	keys, items, err = batch.GetAll()
+	if err != nil {
+		t.Fatalf("unable to get the items from the batch; %v", err)
+	}
+	if len(items) != 1 {
+		t.Fatalf("Expected length of the batch items to be 1 but got %v", len(items))
+	}
+	if len(keys) != 1 {
+		t.Fatalf("Expected length of the batch keys to be 1 but got %v", len(items))
+	}
+	// try adding again after Get.
+	for i := 0; i < int(limit); i++ {
+		if err := batch.Add(i, i); err != nil {
+			t.Fatalf("failed to add item %v to batch after Get; %v", i, err)
+		}
+		if _, ok := batch.GetByKey(i); !ok {
+			t.Fatalf("failed to get the item by key %v after adding", i)
+		}
+	}
+}
+
+func TestBatchWithConcurrency(t *testing.T) {
+	var limit uint32 = 100
+	batch := NewBatch[int, int](limit)
+
+	var wg sync.WaitGroup
+	for i := 0; i < int(limit); i++ {
+		wg.Add(1)
+		go func(item int) {
+			defer wg.Done()
+			if err := batch.Add(item, item); err != nil {
+				t.Errorf("failed to add item %v; %v", item, err)
+				return
+			}
+			if _, ok := batch.GetByKey(item); !ok {
+				t.Errorf("failed to get the item by key %v after adding", item)
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	keys, items, err := batch.GetAll()
+	if err != nil {
+		t.Fatalf("unable to get the items from the batch; %v", err)
+	}
+	if len(items) != int(limit) {
+		t.Fatalf("expected batch length %v but got %v", limit, len(items))
+	}
+	if len(keys) != int(limit) {
+		t.Fatalf("Expected length of the batch keys to be %v but got %v", limit, len(items))
+	}
+	batchLen := batch.Len()
+	if batchLen != 0 {
+		t.Fatalf("expected batch to be empty but still left with %d items", batchLen)
+	}
+}

internal/store/queuestore.go

@@ -167,6 +167,21 @@ func (store *QueueStore[_]) Del(key string) error {
 	return store.del(key)
 }
 
+// DelList - Deletes a list of entries from the store.
+// Returns an error even if one key fails to be deleted.
+func (store *QueueStore[_]) DelList(keys []string) error {
+	store.Lock()
+	defer store.Unlock()
+
+	for _, key := range keys {
+		if err := store.del(key); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // Len returns the entry count.
 func (store *QueueStore[_]) Len() int {
 	store.RLock()

internal/store/store.go

@@ -39,7 +39,7 @@ var ErrNotConnected = errors.New("not connected to target server/service")
 // Target - store target interface
 type Target interface {
 	Name() string
-	SendFromStore(key string) error
+	SendFromStore(key Key) error
 }
 
 // Store - Used to persist items.
@@ -49,16 +49,23 @@ type Store[I any] interface {
 	Len() int
 	List() ([]string, error)
 	Del(key string) error
+	DelList(key []string) error
 	Open() error
 	Extension() string
 }
 
+// Key denotes the key present in the store.
+type Key struct {
+	Name   string
+	IsLast bool
+}
+
 // replayItems - Reads the items from the store and replays.
-func replayItems[I any](store Store[I], doneCh <-chan struct{}, log logger, id string) <-chan string {
-	itemKeyCh := make(chan string)
+func replayItems[I any](store Store[I], doneCh <-chan struct{}, log logger, id string) <-chan Key {
+	keyCh := make(chan Key)
 
 	go func() {
-		defer close(itemKeyCh)
+		defer close(keyCh)
 
 		retryTicker := time.NewTicker(retryInterval)
 		defer retryTicker.Stop()
@@ -68,9 +75,10 @@ func replayItems[I any](store Store[I], doneCh <-chan struct{}, log logger, id s
 			if err != nil {
 				log(context.Background(), fmt.Errorf("store.List() failed with: %w", err), id)
 			} else {
-				for _, name := range names {
+				keyCount := len(names)
+				for i, name := range names {
 					select {
-					case itemKeyCh <- strings.TrimSuffix(name, store.Extension()):
+					case keyCh <- Key{strings.TrimSuffix(name, store.Extension()), keyCount == i+1}:
 					// Get next key.
 					case <-doneCh:
 						return
@@ -86,17 +94,17 @@ func replayItems[I any](store Store[I], doneCh <-chan struct{}, log logger, id s
 		}
 	}()
 
-	return itemKeyCh
+	return keyCh
 }
 
 // sendItems - Reads items from the store and re-plays.
-func sendItems(target Target, itemKeyCh <-chan string, doneCh <-chan struct{}, logger logger) {
+func sendItems(target Target, keyCh <-chan Key, doneCh <-chan struct{}, logger logger) {
 	retryTicker := time.NewTicker(retryInterval)
 	defer retryTicker.Stop()
 
-	send := func(itemKey string) bool {
+	send := func(key Key) bool {
 		for {
-			err := target.SendFromStore(itemKey)
+			err := target.SendFromStore(key)
 			if err == nil {
 				break
 			}
@@ -120,13 +128,13 @@ func sendItems(target Target, itemKeyCh <-chan string, doneCh <-chan struct{}, l
 
 	for {
 		select {
-		case itemKey, ok := <-itemKeyCh:
+		case key, ok := <-keyCh:
 			if !ok {
 				// closed channel.
 				return
 			}
 
-			if !send(itemKey) {
+			if !send(key) {
 				return
 			}
 		case <-doneCh:
@@ -139,8 +147,8 @@ func sendItems(target Target, itemKeyCh <-chan string, doneCh <-chan struct{}, l
 func StreamItems[I any](store Store[I], target Target, doneCh <-chan struct{}, logger logger) {
 	go func() {
 		// Replays the items from the store.
-		itemKeyCh := replayItems(store, doneCh, logger, target.Name())
+		keyCh := replayItems(store, doneCh, logger, target.Name())
 		// Send items from the store.
-		sendItems(target, itemKeyCh, doneCh, logger)
+		sendItems(target, keyCh, doneCh, logger)
 	}()
 }