netsync/chainmgr/storage.go

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