github.com/markusbkk/elvish@v0.0.0-20231204143114-91dc52438621/pkg/store/db_store.go (about)

     1  package store
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  
     8  	bolt "go.etcd.io/bbolt"
     9  	"github.com/markusbkk/elvish/pkg/logutil"
    10  	. "github.com/markusbkk/elvish/pkg/store/storedefs"
    11  )
    12  
    13  var logger = logutil.GetLogger("[store] ")
    14  var initDB = map[string](func(*bolt.Tx) error){}
    15  
    16  // DBStore is the permanent storage backend for elvish. It is not thread-safe.
    17  // In particular, the store may be closed while another goroutine is still
    18  // accessing the store. To prevent bad things from happening, every time the
    19  // main goroutine spawns a new goroutine to operate on the store, it should call
    20  // wg.Add(1) in the main goroutine before spawning another goroutine, and
    21  // call wg.Done() in the spawned goroutine after the operation is finished.
    22  type DBStore interface {
    23  	Store
    24  	Close() error
    25  }
    26  
    27  type dbStore struct {
    28  	db *bolt.DB
    29  	wg sync.WaitGroup // used for registering outstanding operations on the store
    30  }
    31  
    32  func dbWithDefaultOptions(dbname string) (*bolt.DB, error) {
    33  	db, err := bolt.Open(dbname, 0644,
    34  		&bolt.Options{
    35  			Timeout: 1 * time.Second,
    36  		})
    37  	return db, err
    38  }
    39  
    40  // NewStore creates a new Store from the given file.
    41  func NewStore(dbname string) (DBStore, error) {
    42  	db, err := dbWithDefaultOptions(dbname)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  	return NewStoreFromDB(db)
    47  }
    48  
    49  // NewStoreFromDB creates a new Store from a bolt DB.
    50  func NewStoreFromDB(db *bolt.DB) (DBStore, error) {
    51  	logger.Println("initializing store")
    52  	defer logger.Println("initialized store")
    53  	st := &dbStore{
    54  		db: db,
    55  		wg: sync.WaitGroup{},
    56  	}
    57  
    58  	err := db.Update(func(tx *bolt.Tx) error {
    59  		for name, fn := range initDB {
    60  			err := fn(tx)
    61  			if err != nil {
    62  				return fmt.Errorf("failed to %s: %v", name, err)
    63  			}
    64  		}
    65  		return nil
    66  	})
    67  	return st, err
    68  }
    69  
    70  // Close waits for all outstanding operations to finish, and closes the
    71  // database.
    72  func (s *dbStore) Close() error {
    73  	if s == nil || s.db == nil {
    74  		return nil
    75  	}
    76  	s.wg.Wait()
    77  	return s.db.Close()
    78  }