github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/wal/kv/mocked.go (about)

     1  package kv
     2  
     3  import (
     4  	"os"
     5  	"strings"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/alphadose/haxmap"
    10  	"github.com/cockroachdb/errors"
    11  	"github.com/projecteru2/core/types"
    12  )
    13  
    14  // MockedKV .
    15  type MockedKV struct {
    16  	sync.Mutex
    17  	pool    *haxmap.Map[string, []byte]
    18  	nextSeq uint64
    19  }
    20  
    21  // NewMockedKV .
    22  func NewMockedKV() *MockedKV {
    23  	return &MockedKV{
    24  		nextSeq: 1,
    25  		pool:    haxmap.New[string, []byte](),
    26  	}
    27  }
    28  
    29  // Open .
    30  func (m *MockedKV) Open(string, os.FileMode, time.Duration) error {
    31  	return nil
    32  }
    33  
    34  // Close .
    35  func (m *MockedKV) Close() error {
    36  	m.pool.ForEach(func(k string, _ []byte) bool {
    37  		m.pool.Del(k)
    38  		return true
    39  	})
    40  	return nil
    41  }
    42  
    43  // NextSequence .
    44  func (m *MockedKV) NextSequence() (nextSeq uint64, err error) {
    45  	m.Lock()
    46  	defer m.Unlock()
    47  	nextSeq = m.nextSeq
    48  	m.nextSeq++
    49  	return
    50  }
    51  
    52  // Put .
    53  func (m *MockedKV) Put(key, value []byte) (err error) {
    54  	m.pool.Set(string(key), value)
    55  	return
    56  }
    57  
    58  // Get .
    59  func (m *MockedKV) Get(key []byte) (value []byte, err error) {
    60  	value, ok := m.pool.Get(string(key))
    61  	if !ok {
    62  		return value, errors.Wrapf(types.ErrInvaildCount, "no such key: %s", key)
    63  	}
    64  	return
    65  }
    66  
    67  // Delete .
    68  func (m *MockedKV) Delete(key []byte) (err error) {
    69  	m.pool.Del(string(key))
    70  	return
    71  }
    72  
    73  // Scan .
    74  func (m *MockedKV) Scan(prefix []byte) (<-chan ScanEntry, func()) {
    75  	ch := make(chan ScanEntry)
    76  
    77  	exit := make(chan struct{})
    78  	abort := func() {
    79  		close(exit)
    80  	}
    81  
    82  	go func() {
    83  		defer close(ch)
    84  
    85  		dataCh := make(chan MockedScanEntry)
    86  		go func() {
    87  			defer close(dataCh)
    88  			m.pool.ForEach(func(k string, v []byte) bool {
    89  				dataCh <- MockedScanEntry{Key: k, Value: v}
    90  				return true
    91  			})
    92  		}()
    93  
    94  		for {
    95  			select {
    96  			case <-exit:
    97  				return
    98  			case entry, ok := <-dataCh:
    99  				switch {
   100  				case !ok:
   101  					return
   102  				case strings.HasPrefix(entry.Key, string(prefix)):
   103  					ch <- entry
   104  				}
   105  			}
   106  		}
   107  	}()
   108  
   109  	return ch, abort
   110  }
   111  
   112  // MockedScanEntry .
   113  type MockedScanEntry struct {
   114  	Err   error
   115  	Key   string
   116  	Value []byte
   117  }
   118  
   119  // Pair .
   120  func (e MockedScanEntry) Pair() ([]byte, []byte) {
   121  	return []byte(e.Key), e.Value
   122  }
   123  
   124  // Error .
   125  func (e MockedScanEntry) Error() error {
   126  	return e.Err
   127  }