github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/utils/dbutil/switchable/snapshot.go (about)

     1  package switchable
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/unicornultrafoundation/go-u2u/common"
     7  
     8  	"github.com/unicornultrafoundation/go-helios/u2udb"
     9  	"github.com/unicornultrafoundation/go-u2u/utils/dbutil/itergc"
    10  )
    11  
    12  type Snapshot struct {
    13  	u2udb.Snapshot
    14  	mu sync.RWMutex
    15  }
    16  
    17  func (s *Snapshot) SwitchTo(snap u2udb.Snapshot) u2udb.Snapshot {
    18  	s.mu.Lock()
    19  	defer s.mu.Unlock()
    20  	old := s.Snapshot
    21  	s.Snapshot = itergc.Wrap(snap, &sync.Mutex{})
    22  	return old
    23  }
    24  
    25  func Wrap(snap u2udb.Snapshot) *Snapshot {
    26  	s := &Snapshot{}
    27  	s.SwitchTo(snap)
    28  	return s
    29  }
    30  
    31  // Has checks if key is in the exists.
    32  func (s *Snapshot) Has(key []byte) (bool, error) {
    33  	s.mu.RLock()
    34  	defer s.mu.RUnlock()
    35  
    36  	return s.Snapshot.Has(key)
    37  }
    38  
    39  // Get returns key-value pair by key.
    40  func (s *Snapshot) Get(key []byte) ([]byte, error) {
    41  	s.mu.RLock()
    42  	defer s.mu.RUnlock()
    43  
    44  	return s.Snapshot.Get(key)
    45  }
    46  
    47  func (s *Snapshot) Release() {
    48  	s.mu.Lock()
    49  	defer s.mu.Unlock()
    50  
    51  	s.Snapshot.Release()
    52  }
    53  
    54  // NewIterator creates a binary-alphabetical iterator over a subset
    55  // of database content with a particular key prefix, starting at a particular
    56  // initial key (or after, if it does not exist).
    57  func (s *Snapshot) NewIterator(prefix []byte, start []byte) u2udb.Iterator {
    58  	s.mu.RLock()
    59  	defer s.mu.RUnlock()
    60  
    61  	return &switchableIterator{
    62  		mu:       &s.mu,
    63  		upd:      &s.Snapshot,
    64  		cur:      s.Snapshot,
    65  		parentIt: s.Snapshot.NewIterator(prefix, start),
    66  		prefix:   prefix,
    67  		start:    start,
    68  	}
    69  }
    70  
    71  /*
    72   * Iterator
    73   */
    74  
    75  type switchableIterator struct {
    76  	mu       *sync.RWMutex
    77  	upd      *u2udb.Snapshot
    78  	cur      u2udb.Snapshot
    79  	parentIt u2udb.Iterator
    80  
    81  	prefix, start []byte
    82  	key, value    []byte
    83  }
    84  
    85  func (it *switchableIterator) mayReopen() {
    86  	if it.cur != *it.upd {
    87  		// reopen iterator if DB was switched
    88  		it.cur = *it.upd
    89  		if it.key != nil {
    90  			it.start = common.CopyBytes(it.key[len(it.prefix):])
    91  		}
    92  		it.parentIt = it.cur.NewIterator(it.prefix, it.start)
    93  		if it.key != nil {
    94  			_ = it.parentIt.Next() // skip previous key
    95  		}
    96  	}
    97  }
    98  
    99  // Next scans key-value pair by key in lexicographic order. Looks in cache first,
   100  // then - in DB.
   101  func (it *switchableIterator) Next() bool {
   102  	it.mu.Lock()
   103  	defer it.mu.Unlock()
   104  
   105  	it.mayReopen()
   106  
   107  	ok := it.parentIt.Next()
   108  	if !ok {
   109  		it.key = nil
   110  		it.value = nil
   111  		return false
   112  	}
   113  	it.key = it.parentIt.Key()
   114  	it.value = it.parentIt.Value()
   115  	return true
   116  }
   117  
   118  // Error returns any accumulated error. Exhausting all the key/value pairs
   119  // is not considered to be an error. A memory iterator cannot encounter errors.
   120  func (it *switchableIterator) Error() error {
   121  	it.mu.Lock()
   122  	defer it.mu.Unlock()
   123  
   124  	it.mayReopen()
   125  
   126  	return it.parentIt.Error()
   127  }
   128  
   129  // Key returns the key of the current key/value pair, or nil if done. The caller
   130  // should not modify the contents of the returned slice, and its contents may
   131  // change on the next call to Next.
   132  func (it *switchableIterator) Key() []byte {
   133  	return it.key
   134  }
   135  
   136  // Value returns the value of the current key/value pair, or nil if done. The
   137  // caller should not modify the contents of the returned slice, and its contents
   138  // may change on the next call to Next.
   139  func (it *switchableIterator) Value() []byte {
   140  	return it.value
   141  }
   142  
   143  // Release releases associated resources. Release should always succeed and can
   144  // be called multiple times without causing error.
   145  func (it *switchableIterator) Release() {
   146  	it.mu.Lock()
   147  	defer it.mu.Unlock()
   148  
   149  	it.mayReopen()
   150  
   151  	it.parentIt.Release()
   152  }