mirror of
https://github.com/PeernetOfficial/core.git
synced 2026-07-17 02:47:51 +01:00
95 lines
2.2 KiB
Go
95 lines
2.2 KiB
Go
/*
|
|
File Name: Memory.go
|
|
Copyright: 2021 Peernet s.r.o.
|
|
Author: Peter Kleissner
|
|
*/
|
|
|
|
package store
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// MemoryStore is a simple in-memory key-value store for testing purposes.
|
|
type MemoryStore struct {
|
|
mutex *sync.Mutex
|
|
data map[string][]byte
|
|
expireMap map[string]time.Time
|
|
}
|
|
|
|
// NewMemoryStore create a properly initialized memory store.
|
|
func NewMemoryStore() *MemoryStore {
|
|
return &MemoryStore{
|
|
data: make(map[string][]byte),
|
|
mutex: &sync.Mutex{},
|
|
expireMap: make(map[string]time.Time),
|
|
}
|
|
}
|
|
|
|
// ExpireKeys is called to delete all keys that are marked for expiration.
|
|
func (ms *MemoryStore) ExpireKeys() {
|
|
ms.mutex.Lock()
|
|
defer ms.mutex.Unlock()
|
|
for k, v := range ms.expireMap {
|
|
if time.Now().After(v) {
|
|
delete(ms.expireMap, k)
|
|
delete(ms.data, k)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Set stores the key-value pair.
|
|
func (ms *MemoryStore) Set(key []byte, data []byte) error {
|
|
ms.mutex.Lock()
|
|
ms.data[string(key)] = data
|
|
ms.mutex.Unlock()
|
|
return nil
|
|
}
|
|
|
|
// StoreExpire stores the key-value pair and deletes it after the expiration time.
|
|
func (ms *MemoryStore) StoreExpire(key []byte, data []byte, expiration time.Time) error {
|
|
ms.mutex.Lock()
|
|
ms.expireMap[string(key)] = expiration
|
|
ms.data[string(key)] = data
|
|
ms.mutex.Unlock()
|
|
return nil
|
|
}
|
|
|
|
// Get returns the value for the key if present.
|
|
func (ms *MemoryStore) Get(key []byte) (data []byte, found bool) {
|
|
ms.mutex.Lock()
|
|
data, found = ms.data[string(key)]
|
|
ms.mutex.Unlock()
|
|
return data, found
|
|
}
|
|
|
|
// Delete deletes a key-value pair.
|
|
func (ms *MemoryStore) Delete(key []byte) {
|
|
ms.mutex.Lock()
|
|
delete(ms.expireMap, string(key))
|
|
delete(ms.data, string(key))
|
|
ms.mutex.Unlock()
|
|
}
|
|
|
|
// Count returns the count of records stored.
|
|
func (ms *MemoryStore) Count() uint64 {
|
|
ms.mutex.Lock()
|
|
defer ms.mutex.Unlock()
|
|
return uint64(len(ms.data))
|
|
}
|
|
|
|
// Iterate iterates over all records.
|
|
func (ms *MemoryStore) Iterate(callback func(key, value []byte)) {
|
|
ms.mutex.Lock()
|
|
defer ms.mutex.Unlock()
|
|
|
|
for key, value := range ms.data {
|
|
ms.mutex.Unlock() // allow access to the map while calling the callback
|
|
|
|
callback([]byte(key), value)
|
|
|
|
ms.mutex.Lock()
|
|
}
|
|
}
|