sharded_slice.go

package main

import (
	"sync"
	"sync/atomic"
)

type ShardedConcurrentSlice[T any] struct {
	datas [][]T
	mu    []sync.Mutex
	next  atomic.Int64
}

func NewShardedConcurrentSlice[T any](numShards int, prealloc int) *ShardedConcurrentSlice[T] {
	datas := make([][]T, numShards)
	for i := range datas {
		datas[i] = make([]T, 0, prealloc)
	}
	return &ShardedConcurrentSlice[T]{
		datas: datas,
		mu:    make([]sync.Mutex, numShards),
	}
}

func (s *ShardedConcurrentSlice[T]) Append(v T) {
	shard := int(s.next.Add(1)) % len(s.datas)
	s.mu[shard].Lock()
	s.datas[shard] = append(s.datas[shard], v)
	s.mu[shard].Unlock()
}

func (s *ShardedConcurrentSlice[T]) Data() [][]T {
	return s.datas
}