MyFavMirrors
/
minio
mirror of https://github.com/minio/minio.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

Fixes #18124
This commit is contained in:
Praveen raj Mani 2023-10-07 20:37:38 +05:30 committed by GitHub
parent 0de2b9a1b2
commit c27d0583d4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
19 changed files with 438 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
internal/config/notify/parse.go
View File

@ -362,6 +362,10 @@ var (
Key: target.KafkaVersion,
Value: "",
},
config.KV{
Key: target.KafkaBatchSize,
Value: "0",
},
}
)
@ -434,6 +438,15 @@ func GetNotifyKafka(kafkaKVS map[string]config.KVS) (map[string]target.KafkaArgs
versionEnv = versionEnv + config.Default + k
}
batchSizeEnv := target.EnvKafkaBatchSize
if k != config.Default {
batchSizeEnv = batchSizeEnv + config.Default + k
}
batchSize, err := strconv.ParseUint(env.Get(batchSizeEnv, kv.Get(target.KafkaBatchSize)), 10, 32)
if err != nil {
return nil, err
}
kafkaArgs := target.KafkaArgs{
Enable: enabled,
Brokers: brokers,
@ -441,6 +454,7 @@ func GetNotifyKafka(kafkaKVS map[string]config.KVS) (map[string]target.KafkaArgs
QueueDir: env.Get(queueDirEnv, kv.Get(target.KafkaQueueDir)),
QueueLimit: queueLimit,
Version: env.Get(versionEnv, kv.Get(target.KafkaVersion)),
BatchSize: uint32(batchSize),
}
tlsEnableEnv := target.EnvKafkaTLS

6
internal/event/target/amqp.go
View File

@ -291,7 +291,7 @@ func (target *AMQPTarget) Save(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to AMQP091.
func (target *AMQPTarget) SendFromStore(eventKey string) error {
func (target *AMQPTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -302,7 +302,7 @@ func (target *AMQPTarget) SendFromStore(eventKey string) error {
}
defer ch.Close()
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -317,7 +317,7 @@ func (target *AMQPTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - does nothing and available for interface compatibility.

6
internal/event/target/elasticsearch.go
View File

@ -264,7 +264,7 @@ func (target *ElasticsearchTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to Elasticsearch.
func (target *ElasticsearchTarget) SendFromStore(eventKey string) error {
func (target *ElasticsearchTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -277,7 +277,7 @@ func (target *ElasticsearchTarget) SendFromStore(eventKey string) error {
return err
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -295,7 +295,7 @@ func (target *ElasticsearchTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - does nothing and available for interface compatibility.

146
internal/event/target/kafka.go
View File

@ -57,6 +57,7 @@ const (
KafkaClientTLSCert = "client_tls_cert"
KafkaClientTLSKey = "client_tls_key"
KafkaVersion = "version"
KafkaBatchSize = "batch_size"
EnvKafkaEnable = "MINIO_NOTIFY_KAFKA_ENABLE"
EnvKafkaBrokers = "MINIO_NOTIFY_KAFKA_BROKERS"
@ -73,6 +74,9 @@ const (
EnvKafkaClientTLSCert = "MINIO_NOTIFY_KAFKA_CLIENT_TLS_CERT"
EnvKafkaClientTLSKey = "MINIO_NOTIFY_KAFKA_CLIENT_TLS_KEY"
EnvKafkaVersion = "MINIO_NOTIFY_KAFKA_VERSION"
EnvKafkaBatchSize = "MINIO_NOTIFY_KAFKA_BATCH_SIZE"
maxBatchLimit = 100
)
// KafkaArgs - Kafka target arguments.
@ -83,6 +87,7 @@ type KafkaArgs struct {
QueueDir string `json:"queueDir"`
QueueLimit uint64 `json:"queueLimit"`
Version string `json:"version"`
BatchSize uint32 `json:"batchSize"`
TLS struct {
Enable bool `json:"enable"`
RootCAs *x509.CertPool `json:"-"`
@ -122,6 +127,14 @@ func (k KafkaArgs) Validate() error {
return err
}
}
if k.BatchSize > 1 {
if k.QueueDir == "" {
return errors.New("batch should be enabled only if queue dir is enabled")
}
if k.BatchSize > maxBatchLimit {
return fmt.Errorf("batch limit should not exceed %d", maxBatchLimit)
}
}
return nil
}
@ -134,6 +147,7 @@ type KafkaTarget struct {
producer sarama.SyncProducer
config *sarama.Config
store store.Store[event.Event]
batch *store.Batch[string, *sarama.ProducerMessage]
loggerOnce logger.LogOnce
quitCh chan struct{}
}
@ -188,55 +202,33 @@ func (target *KafkaTarget) send(eventData event.Event) error {
if target.producer == nil {
return store.ErrNotConnected
}
objectName, err := url.QueryUnescape(eventData.S3.Object.Key)
msg, err := target.toProducerMessage(eventData)
if err != nil {
return err
}
key := eventData.S3.Bucket.Name + "/" + objectName
data, err := json.Marshal(event.Log{EventName: eventData.EventName, Key: key, Records: []event.Event{eventData}})
if err != nil {
return err
}
msg := sarama.ProducerMessage{
Topic: target.args.Topic,
Key: sarama.StringEncoder(key),
Value: sarama.ByteEncoder(data),
}
_, _, err = target.producer.SendMessage(&msg)
_, _, err = target.producer.SendMessage(msg)
return err
}
// SendFromStore - reads an event from store and sends it to Kafka.
func (target *KafkaTarget) SendFromStore(eventKey string) error {
func (target *KafkaTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
// If batch is enabled, the event will be batched in memory
// and will be committed once the batch is full.
if target.batch != nil {
return target.addToBatch(key)
}
var err error
_, err = target.isActive()
if err != nil {
return err
}
if target.producer == nil {
brokers := []string{}
for _, broker := range target.args.Brokers {
brokers = append(brokers, broker.String())
}
target.producer, err = sarama.NewSyncProducer(brokers, target.config)
if err != nil {
if err != sarama.ErrOutOfBrokers {
return err
}
return store.ErrNotConnected
}
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -248,15 +240,79 @@ func (target *KafkaTarget) SendFromStore(eventKey string) error {
err = target.send(eventData)
if err != nil {
// Sarama opens the ciruit breaker after 3 consecutive connection failures.
if err == sarama.ErrLeaderNotAvailable || err.Error() == "circuit breaker is open" {
if isKafkaConnErr(err) {
return store.ErrNotConnected
}
return err
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
func (target *KafkaTarget) addToBatch(key store.Key) error {
if target.batch.IsFull() {
if err := target.commitBatch(); err != nil {
return err
}
}
if _, ok := target.batch.GetByKey(key.Name); !ok {
eventData, err := target.store.Get(key.Name)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
msg, err := target.toProducerMessage(eventData)
if err != nil {
return err
}
if err = target.batch.Add(key.Name, msg); err != nil {
return err
}
}
// commit the batch if the key is the last one present in the store.
if key.IsLast || target.batch.IsFull() {
return target.commitBatch()
}
return nil
}
func (target *KafkaTarget) commitBatch() error {
if _, err := target.isActive(); err != nil {
return err
}
keys, msgs, err := target.batch.GetAll()
if err != nil {
return err
}
if err = target.producer.SendMessages(msgs); err != nil {
if isKafkaConnErr(err) {
return store.ErrNotConnected
}
return err
}
return target.store.DelList(keys)
}
func (target *KafkaTarget) toProducerMessage(eventData event.Event) (*sarama.ProducerMessage, error) {
objectName, err := url.QueryUnescape(eventData.S3.Object.Key)
if err != nil {
return nil, err
}
key := eventData.S3.Bucket.Name + "/" + objectName
data, err := json.Marshal(event.Log{EventName: eventData.EventName, Key: key, Records: []event.Event{eventData}})
if err != nil {
return nil, err
}
return &sarama.ProducerMessage{
Topic: target.args.Topic,
Key: sarama.StringEncoder(key),
Value: sarama.ByteEncoder(data),
}, nil
}
// Close - closes underneath kafka connection.
@ -335,16 +391,16 @@ func (target *KafkaTarget) initKafka() error {
// These settings are needed to ensure that kafka client doesn't hang on brokers
// refer https://github.com/IBM/sarama/issues/765#issuecomment-254333355
config.Producer.Retry.Max = 2
config.Producer.Retry.Backoff = (10 * time.Second)
config.Producer.Retry.Backoff = (1 * time.Second)
config.Producer.Return.Successes = true
config.Producer.Return.Errors = true
config.Producer.RequiredAcks = 1
config.Producer.Timeout = (10 * time.Second)
config.Net.ReadTimeout = (10 * time.Second)
config.Net.DialTimeout = (10 * time.Second)
config.Net.WriteTimeout = (10 * time.Second)
config.Producer.Timeout = (5 * time.Second)
config.Net.ReadTimeout = (5 * time.Second)
config.Net.DialTimeout = (5 * time.Second)
config.Net.WriteTimeout = (5 * time.Second)
config.Metadata.Retry.Max = 1
config.Metadata.Retry.Backoff = (10 * time.Second)
config.Metadata.Retry.Backoff = (1 * time.Second)
config.Metadata.RefreshFrequency = (15 * time.Minute)
target.config = config
@ -394,8 +450,16 @@ func NewKafkaTarget(id string, args KafkaArgs, loggerOnce logger.LogOnce) (*Kafk
}
if target.store != nil {
if args.BatchSize > 1 {
target.batch = store.NewBatch[string, *sarama.ProducerMessage](args.BatchSize)
}
store.StreamItems(target.store, target, target.quitCh, target.loggerOnce)
}
return target, nil
}
func isKafkaConnErr(err error) bool {
// Sarama opens the ciruit breaker after 3 consecutive connection failures.
return err == sarama.ErrLeaderNotAvailable || err.Error() == "circuit breaker is open"
}

6
internal/event/target/mqtt.go
View File

@ -169,7 +169,7 @@ func (target *MQTTTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to MQTT.
func (target *MQTTTarget) SendFromStore(eventKey string) error {
func (target *MQTTTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -180,7 +180,7 @@ func (target *MQTTTarget) SendFromStore(eventKey string) error {
return err
}
eventData, err := target.store.Get(eventKey)
eventData, err := target.store.Get(key.Name)
if err != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -195,7 +195,7 @@ func (target *MQTTTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Save - saves the events to the store if queuestore is configured, which will

6
internal/event/target/mysql.go
View File

@ -254,7 +254,7 @@ func (target *MySQLTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to MySQL.
func (target *MySQLTarget) SendFromStore(eventKey string) error {
func (target *MySQLTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -273,7 +273,7 @@ func (target *MySQLTarget) SendFromStore(eventKey string) error {
}
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -291,7 +291,7 @@ func (target *MySQLTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - closes underneath connections to MySQL database.

6
internal/event/target/nats.go
View File

@ -335,7 +335,7 @@ func (target *NATSTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to Nats.
func (target *NATSTarget) SendFromStore(eventKey string) error {
func (target *NATSTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -345,7 +345,7 @@ func (target *NATSTarget) SendFromStore(eventKey string) error {
return err
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -359,7 +359,7 @@ func (target *NATSTarget) SendFromStore(eventKey string) error {
return err
}
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - closes underneath connections to NATS server.

6
internal/event/target/nsq.go
View File

@ -178,7 +178,7 @@ func (target *NSQTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to NSQ.
func (target *NSQTarget) SendFromStore(eventKey string) error {
func (target *NSQTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -188,7 +188,7 @@ func (target *NSQTarget) SendFromStore(eventKey string) error {
return err
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -203,7 +203,7 @@ func (target *NSQTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - closes underneath connections to NSQD server.

6
internal/event/target/postgresql.go
View File

@ -251,7 +251,7 @@ func (target *PostgreSQLTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to PostgreSQL.
func (target *PostgreSQLTarget) SendFromStore(eventKey string) error {
func (target *PostgreSQLTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -269,7 +269,7 @@ func (target *PostgreSQLTarget) SendFromStore(eventKey string) error {
}
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and wouldve been already been sent successfully.
@ -287,7 +287,7 @@ func (target *PostgreSQLTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - closes underneath connections to PostgreSQL database.

6
internal/event/target/redis.go
View File

@ -223,7 +223,7 @@ func (target *RedisTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to redis.
func (target *RedisTarget) SendFromStore(eventKey string) error {
func (target *RedisTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
@ -249,7 +249,7 @@ func (target *RedisTarget) SendFromStore(eventKey string) error {
target.firstPing = true
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and would've been already been sent successfully.
@ -267,7 +267,7 @@ func (target *RedisTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - releases the resources used by the pool.

6
internal/event/target/webhook.go
View File

@ -207,12 +207,12 @@ func (target *WebhookTarget) send(eventData event.Event) error {
}
// SendFromStore - reads an event from store and sends it to webhook.
func (target *WebhookTarget) SendFromStore(eventKey string) error {
func (target *WebhookTarget) SendFromStore(key store.Key) error {
if err := target.init(); err != nil {
return err
}
eventData, eErr := target.store.Get(eventKey)
eventData, eErr := target.store.Get(key.Name)
if eErr != nil {
// The last event key in a successful batch will be sent in the channel atmost once by the replayEvents()
// Such events will not exist and would've been already been sent successfully.
@ -230,7 +230,7 @@ func (target *WebhookTarget) SendFromStore(eventKey string) error {
}
// Delete the event from store.
return target.store.Del(eventKey)
return target.store.Del(key.Name)
}
// Close - does nothing and available for interface compatibility.

4
internal/event/targetlist.go
View File

@ -22,6 +22,8 @@ import (
"strings"
"sync"
"sync/atomic"
"github.com/minio/minio/internal/store"
)
const (
@ -34,7 +36,7 @@ type Target interface {
ID() TargetID
IsActive() (bool, error)
Save(Event) error
SendFromStore(string) error
SendFromStore(store.Key) error
Close() error
Store() TargetStore
}

4
internal/event/targetlist_test.go
View File

@ -23,6 +23,8 @@ import (
"reflect"
"testing"
"time"
"github.com/minio/minio/internal/store"
)
type ExampleTarget struct {
@ -61,7 +63,7 @@ func (target ExampleTarget) send(eventData Event) error {
}
// SendFromStore - interface compatible method does no-op.
func (target *ExampleTarget) SendFromStore(eventKey string) error {
func (target *ExampleTarget) SendFromStore(_ store.Key) error {
return nil
}

6
internal/logger/target/http/http.go
View File

@ -349,9 +349,9 @@ func New(config Config) *Target {
}
// SendFromStore - reads the log from store and sends it to webhook.
func (h *Target) SendFromStore(key string) (err error) {
func (h *Target) SendFromStore(key store.Key) (err error) {
var eventData interface{}
eventData, err = h.store.Get(key)
eventData, err = h.store.Get(key.Name)
if err != nil {
if os.IsNotExist(err) {
return nil
@ -372,7 +372,7 @@ func (h *Target) SendFromStore(key string) (err error) {
return err
}
// Delete the event from store.
return h.store.Del(key)
return h.store.Del(key.Name)
}
// Send the log message 'entry' to the http target.

6
internal/logger/target/kafka/kafka.go
View File

@ -361,8 +361,8 @@ func (h *Target) Send(ctx context.Context, entry interface{}) error {
}
// SendFromStore - reads the log from store and sends it to kafka.
func (h *Target) SendFromStore(key string) (err error) {
auditEntry, err := h.store.Get(key)
func (h *Target) SendFromStore(key store.Key) (err error) {
auditEntry, err := h.store.Get(key.Name)
if err != nil {
if os.IsNotExist(err) {
return nil
@ -376,7 +376,7 @@ func (h *Target) SendFromStore(key string) (err error) {
return
}
// Delete the event from store.
return h.store.Del(key)
return h.store.Del(key.Name)
}
// Cancel - cancels the target

117
internal/store/batch.go Normal file
View File

@ -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,
}
}

126
internal/store/batch_test.go Normal file
View File

@ -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)
}
}

15
internal/store/queuestore.go
View File

@ -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()

36
internal/store/store.go
View File

@ -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)
}()
}