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

     1  package dbcounter
     2  
     3  import (
     4  	"fmt"
     5  	"sync/atomic"
     6  
     7  	"github.com/unicornultrafoundation/go-u2u/log"
     8  
     9  	"github.com/unicornultrafoundation/go-helios/u2udb"
    10  )
    11  
    12  type DBProducer struct {
    13  	u2udb.IterableDBProducer
    14  	warn bool
    15  }
    16  
    17  type Iterator struct {
    18  	u2udb.Iterator
    19  	itCounter *int64
    20  	start     []byte
    21  	prefix    []byte
    22  }
    23  
    24  type Snapshot struct {
    25  	u2udb.Snapshot
    26  	snCounter *int64
    27  }
    28  
    29  type Store struct {
    30  	u2udb.Store
    31  	name      string
    32  	snCounter int64
    33  	itCounter int64
    34  	warn      bool
    35  }
    36  
    37  func Wrap(db u2udb.IterableDBProducer, warn bool) u2udb.IterableDBProducer {
    38  	return &DBProducer{db, warn}
    39  }
    40  
    41  func WrapStore(s u2udb.Store, name string, warn bool) *Store {
    42  	return &Store{
    43  		Store: s,
    44  		name:  name,
    45  		warn:  warn,
    46  	}
    47  }
    48  
    49  func (ds *Store) Close() error {
    50  	itCounter, snCounter := atomic.LoadInt64(&ds.itCounter), atomic.LoadInt64(&ds.snCounter)
    51  	if itCounter != 0 || snCounter != 0 {
    52  		err := fmt.Errorf("%s DB leak: %d iterators, %d snapshots", ds.name, itCounter, snCounter)
    53  		if ds.warn {
    54  			log.Warn("Possible " + err.Error())
    55  		} else {
    56  			return err
    57  		}
    58  	}
    59  	return ds.Store.Close()
    60  }
    61  
    62  func (ds *Snapshot) Release() {
    63  	atomic.AddInt64(ds.snCounter, -1)
    64  	ds.Snapshot.Release()
    65  }
    66  
    67  func (ds *Store) NewIterator(prefix []byte, start []byte) u2udb.Iterator {
    68  	atomic.AddInt64(&ds.itCounter, 1)
    69  	return &Iterator{
    70  		Iterator:  ds.Store.NewIterator(prefix, start),
    71  		itCounter: &ds.itCounter,
    72  		start:     start,
    73  		prefix:    prefix,
    74  	}
    75  }
    76  
    77  func (it *Iterator) Release() {
    78  	atomic.AddInt64(it.itCounter, -1)
    79  	it.Iterator.Release()
    80  }
    81  
    82  func (ds *Store) GetSnapshot() (u2udb.Snapshot, error) {
    83  	atomic.AddInt64(&ds.snCounter, 1)
    84  	snapshot, err := ds.Store.GetSnapshot()
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	return &Snapshot{
    89  		Snapshot:  snapshot,
    90  		snCounter: &ds.snCounter,
    91  	}, nil
    92  }
    93  
    94  func (db *DBProducer) OpenDB(name string) (u2udb.Store, error) {
    95  	s, err := db.IterableDBProducer.OpenDB(name)
    96  	if err != nil {
    97  		return nil, err
    98  	}
    99  	return WrapStore(s, name, db.warn), nil
   100  }