github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/contract/trace_store.go (about)

     1  package contract
     2  
     3  import (
     4  	"encoding/json"
     5  
     6  	log "github.com/sirupsen/logrus"
     7  
     8  	dbm "github.com/bytom/bytom/database/leveldb"
     9  	"github.com/bytom/bytom/protocol/bc"
    10  )
    11  
    12  const (
    13  	colon = byte(0x3a)
    14  
    15  	instance byte = iota + 1
    16  	chainStatus
    17  )
    18  
    19  var (
    20  	instancePrefixKey    = []byte{instance, colon}
    21  	chainStatusPrefixKey = []byte{chainStatus, colon}
    22  )
    23  
    24  func instanceKey(traceID string) []byte {
    25  	return append(instancePrefixKey, []byte(traceID)...)
    26  }
    27  
    28  func chainStatusKey() []byte {
    29  	return chainStatusPrefixKey
    30  }
    31  
    32  
    33  type TraceStore struct {
    34  	db dbm.DB
    35  }
    36  
    37  func NewTraceStore(db dbm.DB) *TraceStore {
    38  	return &TraceStore{db: db}
    39  }
    40  
    41  // GetInstance return instance by given trace id
    42  func (t *TraceStore) GetInstance(traceID string) (*Instance, error) {
    43  	key := instanceKey(traceID)
    44  	data := t.db.Get(key)
    45  	instance := &Instance{}
    46  	if err := json.Unmarshal(data, instance); err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	return instance, nil
    51  }
    52  
    53  // LoadInstances used to load all instances in db
    54  func (t *TraceStore) LoadInstances() ([]*Instance, error) {
    55  	iter := t.db.IteratorPrefix(instancePrefixKey)
    56  	defer iter.Release()
    57  
    58  	var instances []*Instance
    59  	for iter.Next() {
    60  		instance := &Instance{}
    61  		if err := json.Unmarshal(iter.Value(), instance); err != nil {
    62  			return nil, err
    63  		}
    64  
    65  		instances = append(instances, instance)
    66  	}
    67  	return instances, nil
    68  }
    69  
    70  // SaveInstances used to batch save multiple instances
    71  func (t *TraceStore) SaveInstances(instances []*Instance) error {
    72  	batch := t.db.NewBatch()
    73  	if err := t.saveInstances(instances, batch); err != nil {
    74  		return err
    75  	}
    76  
    77  	batch.Write()
    78  	return nil
    79  }
    80  
    81  // RemoveInstance delete a instance by given trace id
    82  func (t *TraceStore) RemoveInstance(traceID string) {
    83  	key := instanceKey(traceID)
    84  	t.db.Delete(key)
    85  }
    86  
    87  // SaveInstancesWithStatus batch save the instances and chain status
    88  func (t *TraceStore) SaveInstancesWithStatus(instances []*Instance, blockHeight uint64, blockHash bc.Hash) error {
    89  	batch := t.db.NewBatch()
    90  	if err := t.saveInstances(instances, batch); err != nil {
    91  		return err
    92  	}
    93  
    94  	chainData, err := json.Marshal(&ChainStatus{BlockHeight: blockHeight, BlockHash: blockHash})
    95  	if err != nil {
    96  		return err
    97  	}
    98  
    99  	batch.Set(chainStatusKey(), chainData)
   100  	batch.Write()
   101  	return nil
   102  }
   103  
   104  // GetChainStatus return the current chain status
   105  func (t *TraceStore) GetChainStatus() *ChainStatus {
   106  	data := t.db.Get(chainStatusKey())
   107  	if data == nil {
   108  		return nil
   109  	}
   110  
   111  	chainStatus := &ChainStatus{}
   112  	if err := json.Unmarshal(data, chainStatus); err != nil {
   113  		log.WithFields(log.Fields{"module": logModule, "err": err}).Fatal("get chain status from trace store")
   114  	}
   115  
   116  	return chainStatus
   117  }
   118  
   119  // SaveChainStatus save the chain status
   120  func (t *TraceStore) SaveChainStatus(chainStatus *ChainStatus) error {
   121  	data, err := json.Marshal(chainStatus)
   122  	if err != nil {
   123  		return err
   124  	}
   125  
   126  	t.db.Set(chainStatusKey(), data)
   127  	return nil
   128  }
   129  
   130  func (t *TraceStore) saveInstances(instances []*Instance, batch dbm.Batch) error {
   131  	for _, inst := range instances {
   132  		key := instanceKey(inst.TraceID)
   133  		data, err := json.Marshal(inst)
   134  		if err != nil {
   135  			return err
   136  		}
   137  
   138  		batch.Set(key, data)
   139  	}
   140  	return nil
   141  }