storage/leveldb/leveldb.go

package leveldb

import (
	"encoding/json"
	"time"

	"github.com/syndtr/goleveldb/leveldb"
)

// data structure for storing items in the database
type item struct {
	Value    []byte    `json:"value"`
	ExpireAt time.Time `json:"expire_at"`
}

// Storage interface that is implemented by storage providers
type Storage struct {
	db         *leveldb.DB
	gcInterval time.Duration
	done       chan struct{}
}

// New creates a new memory storage
func New(config ...Config) *Storage {
	cfg := configDefault(config...)

	db, err := leveldb.OpenFile(cfg.Path, nil)
	if err != nil {
		panic(err)
	}

	store := &Storage{
		db:         db,
		gcInterval: cfg.GCInterval,
		done:       make(chan struct{}),
	}

	go store.gc()

	return store
}

// Get value by key
func (s *Storage) Get(key []byte) ([]byte, error) {
	if len(key) <= 0 {
		return nil, nil
	}

	data, err := s.db.Get(key, nil)
	if err != nil {
		return nil, nil
	}

	var stored item
	if err := json.Unmarshal(data, &stored); err != nil {
		return data, nil
	}

	if !stored.ExpireAt.IsZero() && time.Now().After(stored.ExpireAt) {
		if err := s.Delete(string(key)); err != nil {
			return nil, err
		}
		return nil, nil
	}

	return stored.Value, nil
}

// Set key with value
func (s *Storage) Set(key, value []byte, exp time.Duration) error {
	if len(key) <= 0 || len(value) <= 0 {
		return nil
	}
	if exp == 0 {
		return s.db.Put(key, value, nil)
	}

	data := item{
		Value:    value,
		ExpireAt: time.Now().Add(exp),
	}

	encoded, err := json.Marshal(data)
	if err != nil {
		return err
	}
	return s.db.Put(key, encoded, nil)
}

// Delete key by key
func (s *Storage) Delete(key string) error {
	if len(key) <= 0 {
		return nil
	}

	return s.db.Delete([]byte(key), nil)
}

// Reset all keys
func (s *Storage) Reset() error {
	iter := s.db.NewIterator(nil, nil)
	defer iter.Release()

	for iter.Next() {
		key := iter.Key()
		if err := s.db.Delete(key, nil); err != nil {
			return err
		}
	}

	return iter.Error()
}

// Close the memory storage
func (s *Storage) Close() error {
	s.done <- struct{}{} // GC stop
	close(s.done)
	return s.db.Close()
}

// Return database client
func (s *Storage) Conn() *leveldb.DB {
	return s.db
}

// gc is a helper function to clean up expired keys
func (s *Storage) gc() {
	ticker := time.NewTicker(s.gcInterval)
	defer ticker.Stop()

	for {
		select {
		case <-s.done:
			return
		case <-ticker.C:
			iter := s.db.NewIterator(nil, nil)
			batch := new(leveldb.Batch)

			for iter.Next() {
				key := iter.Key()
				data := iter.Value()

				var stored item
				if err := json.Unmarshal(data, &stored); err != nil {
					continue
				}
				if !stored.ExpireAt.IsZero() && time.Now().After(stored.ExpireAt) {
					batch.Delete(key)
				}
			}

			iter.Release()

			if batch.Len() > 0 {
				_ = s.db.Write(batch, nil)
			}
		}
	}
}