github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/netsync/chainmgr/storage.go (about)

     1  package chainmgr
     2  
     3  import (
     4  	"encoding/binary"
     5  	"sync"
     6  
     7  	dbm "github.com/bytom/bytom/database/leveldb"
     8  	"github.com/bytom/bytom/errors"
     9  	"github.com/bytom/bytom/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  	if err := block.UnmarshalText(binaryBlock); err != nil {
   162  		return nil, err
   163  	}
   164  
   165  	return block, nil
   166  }