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

     1  package asyncflushproducer
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  
     7  	"github.com/unicornultrafoundation/go-helios/u2udb"
     8  	"github.com/unicornultrafoundation/go-u2u/metrics"
     9  )
    10  
    11  type Producer struct {
    12  	u2udb.FullDBProducer
    13  	mu    sync.Mutex
    14  	dbs   map[string]*store
    15  	stats metrics.Meter
    16  
    17  	threshold uint64
    18  }
    19  
    20  func Wrap(backend u2udb.FullDBProducer, threshold uint64) *Producer {
    21  	return &Producer{
    22  		stats:          metrics.NewMeterForced(),
    23  		FullDBProducer: backend,
    24  		dbs:            make(map[string]*store),
    25  		threshold:      threshold,
    26  	}
    27  }
    28  
    29  func (f *Producer) OpenDB(name string) (u2udb.Store, error) {
    30  	f.mu.Lock()
    31  	defer f.mu.Unlock()
    32  	// open existing DB
    33  	openedDB := f.dbs[name]
    34  	if openedDB != nil {
    35  		return openedDB, nil
    36  	}
    37  	// create new DB
    38  	db, err := f.FullDBProducer.OpenDB(name)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	if f.dbs[name] != nil {
    43  		return nil, errors.New("already opened")
    44  	}
    45  	wrapped := &store{
    46  		Store: db,
    47  		CloseFn: func() error {
    48  			f.mu.Lock()
    49  			delete(f.dbs, name)
    50  			f.mu.Unlock()
    51  			return db.Close()
    52  		},
    53  	}
    54  	f.dbs[name] = wrapped
    55  	return wrapped, nil
    56  }
    57  
    58  func (f *Producer) Flush(id []byte) error {
    59  	f.stats.Mark(int64(f.FullDBProducer.NotFlushedSizeEst()))
    60  
    61  	err := f.FullDBProducer.Flush(id)
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	// trigger flushing data to disk if throughput is below a threshold
    67  	if uint64(f.stats.Rate1()) <= f.threshold {
    68  		go func() {
    69  			f.mu.Lock()
    70  			defer f.mu.Unlock()
    71  			for _, db := range f.dbs {
    72  				_, _ = db.Stat("async_flush")
    73  			}
    74  		}()
    75  	}
    76  
    77  	return nil
    78  }