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

     1  package itergc
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/unicornultrafoundation/go-helios/u2udb"
     7  )
     8  
     9  type Snapshot struct {
    10  	u2udb.Snapshot
    11  	nextID uint64
    12  	iters  map[uint64]u2udb.Iterator
    13  	mu     sync.Locker
    14  }
    15  
    16  type Iterator struct {
    17  	u2udb.Iterator
    18  	mu    sync.Locker
    19  	id    uint64
    20  	iters map[uint64]u2udb.Iterator
    21  }
    22  
    23  // Wrap snapshot to automatically close all pending iterators upon snapshot release
    24  func Wrap(snapshot u2udb.Snapshot, mu sync.Locker) *Snapshot {
    25  	return &Snapshot{
    26  		Snapshot: snapshot,
    27  		iters:    make(map[uint64]u2udb.Iterator),
    28  		mu:       mu,
    29  	}
    30  }
    31  
    32  func (s *Snapshot) NewIterator(prefix []byte, start []byte) u2udb.Iterator {
    33  	s.mu.Lock()
    34  	defer s.mu.Unlock()
    35  	it := s.Snapshot.NewIterator(prefix, start)
    36  	id := s.nextID
    37  	s.iters[id] = it
    38  	s.nextID++
    39  
    40  	return &Iterator{
    41  		Iterator: it,
    42  		mu:       s.mu,
    43  		id:       id,
    44  		iters:    s.iters,
    45  	}
    46  }
    47  
    48  func (s *Iterator) Release() {
    49  	s.mu.Lock()
    50  	defer s.mu.Unlock()
    51  	s.Iterator.Release()
    52  	delete(s.iters, s.id)
    53  }
    54  
    55  func (s *Snapshot) Release() {
    56  	s.mu.Lock()
    57  	defer s.mu.Unlock()
    58  	// release all pending iterators
    59  	for _, it := range s.iters {
    60  		it.Release()
    61  	}
    62  	s.iters = nil
    63  	s.Snapshot.Release()
    64  }