diff options
Diffstat (limited to 'pkg/core/shards.go')
| -rw-r--r-- | pkg/core/shards.go | 134 |
1 files changed, 134 insertions, 0 deletions
diff --git a/pkg/core/shards.go b/pkg/core/shards.go new file mode 100644 index 0000000..5ede61c --- /dev/null +++ b/pkg/core/shards.go @@ -0,0 +1,134 @@ +package core + +import ( + "hash/fnv" + "sort" + "sync" +) +const NumShards = 16 + +type Shard struct{ + Mu sync.RWMutex + Id int + Data map[string]Item +} + +func (db *RedisDB) GetShard(key string)*Shard{ + hash := fnv.New64a() + hash.Write([]byte(key)) + val := hash.Sum64() + + return db.Shards[val%NumShards] +} + +func (db *RedisDB) Execute(key string, fn func(*Shard) interface{}) interface{} { + shard := db.GetShard(key) + shard.Mu.Lock() + defer shard.Mu.Unlock() + return fn(shard) +} +func (db *RedisDB) ExecuteRead(key string, fn func(*Shard) interface{}) interface{} { + shard := db.GetShard(key) + shard.Mu.RLock() + defer shard.Mu.RUnlock() + return fn(shard) +} +func ( db *RedisDB) ExecuteMulti(keys []string, fn func([]*Shard) interface{})interface{}{ + if len(keys) == 1{ + shard := db.GetShard(keys[0]) + shard.Mu.Lock() + defer shard.Mu.Unlock() + Shards := []*Shard{shard} + return fn(Shards) + }else{ + shardMap := make(map[int]*Shard) + for _, k := range keys { + shard := db.GetShard(k) + shardMap[shard.Id] = shard + } + var sortedIDs []int + for Id := range shardMap { + sortedIDs = append(sortedIDs, Id) + } + sort.Ints(sortedIDs) + for _, Id := range sortedIDs{ + shardMap[Id].Mu.Lock() + } + defer func(){ + for i := len(sortedIDs)-1; i >= 0; i --{ + shardMap[sortedIDs[i]].Mu.Unlock() + } + }() + + Shards := make([]*Shard, 0, len(shardMap)) + for _, Id := range sortedIDs { + Shards = append(Shards, shardMap[Id]) + } + return fn(Shards) + } +} + +func ( db *RedisDB) ExecuteAll(fn func([]*Shard) interface{})interface{}{ + for _, shard := range db.Shards { + shard.Mu.Lock() + } + + defer func(){ + for i := len(db.Shards)-1;i >= 0; i --{ + db.Shards[i].Mu.Unlock() + } + }() + + return fn(db.Shards) +} + +func ( db *RedisDB) ExecuteReadAll(fn func([]*Shard) interface{})interface{}{ + for _, shard := range db.Shards { + shard.Mu.RLock() + } + + defer func(){ + for i := len(db.Shards)-1;i >= 0; i --{ + db.Shards[i].Mu.RUnlock() + } + }() + + return fn(db.Shards) +} + +func ( db *RedisDB) Lock(keys []string){ + if len(keys) == 1{ + shard := db.GetShard(keys[0]) + shard.Mu.Lock() + }else{ + shardMap := make(map[int]*Shard) + for _, k := range keys { + shard := db.GetShard(k) + shardMap[shard.Id] = shard + } + var sortedIDs []int + for Id := range shardMap { + sortedIDs = append(sortedIDs, Id) + } + sort.Ints(sortedIDs) + for _, Id := range sortedIDs{ + shardMap[Id].Mu.Lock() + } + } +} +func ( db *RedisDB) Unlock(keys []string){ + shardMap := make(map[int]*Shard) + for _, k := range keys { + shard := db.GetShard(k) + shardMap[shard.Id] = shard + } + var sortedIDs []int + for Id := range shardMap { + sortedIDs = append(sortedIDs, Id) + } + sort.Ints(sortedIDs) + + for i := len(sortedIDs)-1; i >= 0; i --{ + shardMap[sortedIDs[i]].Mu.Unlock() + } +}
\ No newline at end of file |