netsync/chainmgr/storage.go

   1 package chainmgr
   2
   3 import (
   4         "encoding/binary"
   5         "sync"
   6
   7         dbm "github.com/bytom/vapor/database/leveldb"
   8         "github.com/bytom/vapor/errors"
   9         "github.com/bytom/vapor/protocol/bc/types"
  10 )
  11
  12 var (
  13         maxByteOfStorageRAM = 800 * 1024 * 1024 //100MB
  14         errStorageFindBlock = errors.New("can't find block from storage")
  15         errDBFindBlock      = errors.New("can't find block from DB")
  16 )
  17
  18 // LocalStore is the interface for persistent storage
  19 type LocalStore interface {
  20         writeBlock(block *types.Block) error
  21         readBlock(height uint64) (*types.Block, error)
  22         clearData()
  23 }
  24
  25 type blockStorage struct {
  26         block  *types.Block
  27         peerID string
  28         size   int
  29         isRAM  bool
  30 }
  31
  32 type storage struct {
  33         actualUsage int
  34         blocks      map[uint64]*blockStorage
  35         localStore  LocalStore
  36         mux         sync.RWMutex
  37 }
  38
  39 func newStorage(db dbm.DB) *storage {
  40         DBStorage := newDBStore(db)
  41         DBStorage.clearData()
  42         return &storage{
  43                 blocks:     make(map[uint64]*blockStorage),
  44                 localStore: DBStorage,
  45         }
  46 }
  47
  48 func (s *storage) writeBlocks(peerID string, blocks []*types.Block) error {
  49         s.mux.Lock()
  50         defer s.mux.Unlock()
  51
  52         for _, block := range blocks {
  53                 binaryBlock, err := block.MarshalText()
  54                 if err != nil {
  55                         return errors.Wrap(err, "Marshal block header")
  56                 }
  57
  58                 if len(binaryBlock)+s.actualUsage < maxByteOfStorageRAM {
  59                         s.blocks[block.Height] = &blockStorage{block: block, peerID: peerID, size: len(binaryBlock), isRAM: true}
  60                         s.actualUsage += len(binaryBlock)
  61                         continue
  62                 }
  63
  64                 if err := s.localStore.writeBlock(block); err != nil {
  65                         return err
  66                 }
  67
  68                 s.blocks[block.Height] = &blockStorage{peerID: peerID, isRAM: false}
  69         }
  70
  71         return nil
  72 }
  73
  74 func (s *storage) readBlock(height uint64) (*blockStorage, error) {
  75         s.mux.RLock()
  76         defer s.mux.RUnlock()
  77
  78         blockStore, ok := s.blocks[height]
  79         if !ok {
  80                 return nil, errStorageFindBlock
  81         }
  82
  83         if blockStore.isRAM {
  84                 return blockStore, nil
  85         }
  86
  87         block, err := s.localStore.readBlock(height)
  88         if err != nil {
  89                 return nil, err
  90         }
  91
  92         blockStore.block = block
  93         return blockStore, nil
  94 }
  95
  96 // deleteBlock delete blocks in memory
  97 func (s *storage) deleteBlock(height uint64) {
  98         s.mux.RLock()
  99         defer s.mux.RUnlock()
 100
 101         blockStore, ok := s.blocks[height]
 102         if !ok {
 103                 return
 104         }
 105
 106         if blockStore.isRAM {
 107                 s.actualUsage -= blockStore.size
 108                 delete(s.blocks, height)
 109         }
 110 }
 111
 112 func (s *storage) resetParameter() {
 113         s.mux.Lock()
 114         defer s.mux.Unlock()
 115
 116         s.blocks = make(map[uint64]*blockStorage)
 117         s.actualUsage = 0
 118         s.localStore.clearData()
 119 }
 120
 121 type levelDBStorage struct {
 122         db dbm.DB
 123 }
 124
 125 func newDBStore(db dbm.DB) *levelDBStorage {
 126         return &levelDBStorage{
 127                 db: db,
 128         }
 129 }
 130
 131 func (ls *levelDBStorage) clearData() {
 132         iter := ls.db.Iterator()
 133         defer iter.Release()
 134
 135         for iter.Next() {
 136                 ls.db.Delete(iter.Key())
 137         }
 138 }
 139
 140 func (ls *levelDBStorage) writeBlock(block *types.Block) error {
 141         binaryBlock, err := block.MarshalText()
 142         if err != nil {
 143                 return err
 144         }
 145
 146         key := make([]byte, 8)
 147         binary.BigEndian.PutUint64(key, block.Height)
 148         ls.db.Set(key, binaryBlock)
 149         return nil
 150 }
 151
 152 func (ls *levelDBStorage) readBlock(height uint64) (*types.Block, error) {
 153         key := make([]byte, 8)
 154         binary.BigEndian.PutUint64(key, height)
 155         binaryBlock := ls.db.Get(key)
 156         if binaryBlock == nil {
 157                 return nil, errDBFindBlock
 158         }
 159
 160         block := &types.Block{}
 161         return block, block.UnmarshalText(binaryBlock)
 162 }