protocol/prottest/memstore/memstore.go

   1 // MemStore is a Store implementation that
   2 // keeps all blockchain state in memory.
   3 //
   4 // It is used in tests to avoid needing a database.
   5 package memstore
   6
   7 import (
   8         "context"
   9         "fmt"
  10         "sync"
  11
  12         "github.com/bytom/protocol/bc/legacy"
  13         "github.com/bytom/protocol/state"
  14 )
  15
  16 // MemStore satisfies the Store interface.
  17 type MemStore struct {
  18         mu          sync.Mutex
  19         Blocks      map[uint64]*legacy.Block
  20         State       *state.Snapshot
  21         StateHeight uint64
  22 }
  23
  24 // New returns a new MemStore
  25 func New() *MemStore {
  26         return &MemStore{Blocks: make(map[uint64]*legacy.Block)}
  27 }
  28
  29 func (m *MemStore) Height() uint64 {
  30         m.mu.Lock()
  31         defer m.mu.Unlock()
  32
  33         return uint64(len(m.Blocks))
  34
  35 }
  36
  37 func (m *MemStore) SaveBlock(b *legacy.Block) error {
  38         m.mu.Lock()
  39         defer m.mu.Unlock()
  40
  41         existing, ok := m.Blocks[b.Height]
  42         if ok && existing.Hash() != b.Hash() {
  43                 return fmt.Errorf("already have a block at height %d", b.Height)
  44         }
  45         m.Blocks[b.Height] = b
  46         return nil
  47 }
  48
  49 func (m *MemStore) SaveSnapshot(ctx context.Context, height uint64, snapshot *state.Snapshot) error {
  50         m.mu.Lock()
  51         defer m.mu.Unlock()
  52
  53         m.State = state.Copy(snapshot)
  54         m.StateHeight = height
  55         return nil
  56 }
  57
  58 func (m *MemStore) GetBlock(height uint64) (*legacy.Block, error) {
  59         m.mu.Lock()
  60         defer m.mu.Unlock()
  61         b, ok := m.Blocks[height]
  62         if !ok {
  63                 return nil, fmt.Errorf("memstore: no block at height %d", height)
  64         }
  65         return b, nil
  66 }
  67
  68 func (m *MemStore) LatestSnapshot(context.Context) (*state.Snapshot, uint64, error) {
  69         m.mu.Lock()
  70         defer m.mu.Unlock()
  71
  72         if m.State == nil {
  73                 m.State = state.Empty()
  74         }
  75         return state.Copy(m.State), m.StateHeight, nil
  76 }
  77
  78 func (m *MemStore) FinalizeBlock(context.Context, uint64) error { return nil }