Make MutexMap shards configurable

storage/badger_storage.go

@@ -80,7 +80,8 @@ type BadgerStorage struct {
 	encoder  *Encoder
 	compress bool
 
-	writer *utils.MutexMap
+	writer       *utils.MutexMap
+	writerShards int
 
 	// Track the closed status to ensure we exit garbage
 	// collection when the db closes.
@@ -197,12 +198,16 @@ func NewBadgerStorage(
 		closed:        make(chan struct{}),
 		pool:          NewBufferPool(),
 		compress:      true,
-		writer:        utils.NewMutexMap(utils.DefaultShards),
+		writerShards:  utils.DefaultShards,
 	}
 	for _, opt := range storageOptions {
 		opt(b)
 	}
 
+	// Initialize utis.MutexMap used to track granular
+	// write transactions.
+	b.writer = utils.NewMutexMap(b.writerShards)
+
 	db, err := badger.Open(b.badgerOptions)
 	if err != nil {
 		return nil, fmt.Errorf("%w: %v", ErrDatabaseOpenFailed, err)

storage/badger_storage_configuration.go

@@ -56,3 +56,13 @@ func WithCustomSettings(settings badger.Options) BadgerOption {
 		b.badgerOptions = settings
 	}
 }
+
+// WithWriterShards overrides the default shards used
+// in the writer utils.MutexMap. It is recommended
+// to set this value to your write concurrency to prevent
+// lock contention.
+func WithWriterShards(shards int) BadgerOption {
+	return func(b *BadgerStorage) {
+		b.writerShards = shards
+	}
+}