Merge pull request #800 from molecula/rbf-exclusive-write

Add RBF exclusive lock mode
This commit is contained in:
Ben Johnson 2020-09-04 13:41:52 -06:00 • committed by GitHub
commit 278d518f03
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 241 additions and 75 deletions

164
rbf/db.go
View file

@ -47,8 +47,9 @@ type DB struct {
txs map[*Tx]struct{} // active transactions
opened bool // true if open
mu sync.RWMutex // general mutex
rwmu sync.Mutex // mutex for restricting single writer
mu sync.RWMutex // general mutex
rwmu sync.Mutex // mutex for restricting single writer
exclmu sync.RWMutex // mutex for locking out everyone but a single writer
// Path represents the path to the database file.
Path string
@ -134,7 +135,7 @@ func (db *DB) Open() (err error) {
// Open write-ahead log & checkpoint to the end since no transactions are open.
if err := db.openWALSegments(); err != nil {
return fmt.Errorf("wal open: %w", err)
} else if err := db.checkpoint(); err != nil {
} else if err := db.checkpoint(true); err != nil {
return fmt.Errorf("checkpoint: %w", err)
}
@ -142,7 +143,6 @@ func (db *DB) Open() (err error) {
}
func (db *DB) openWALSegments() error {
fis, err := ioutil.ReadDir(db.WALPath())
if err != nil {
return fmt.Errorf("read dir: %w", err)
@ -162,11 +162,11 @@ func (db *DB) openWALSegments() error {
db.segments = append(db.segments, segment)
}
// Truncate last WAL page if it is a bitmap header.
if segment := db.activeWALSegment(); segment != nil {
if err := segment.trimBitmapHeaderTrailer(); err != nil {
return err
}
// Truncate everything after the last successful meta page.
if walID, err := db.findLastWALMetaPage(); err != nil {
return err
} else if err := db.truncateWALAfter(walID); err != nil {
return err
}
return nil
@ -175,8 +175,9 @@ func (db *DB) openWALSegments() error {
// checkpoint copies pages from WAL segments into the main DB file. This can
// only copy pages that aren't in use by an active transaction. The page map
// is rebuilt as well for all WAL pages still in use.
func (db *DB) checkpoint() error {
//
// If exclusive is true, all WAL writes are flushed to disk.
func (db *DB) checkpoint(exclusive bool) error {
if !db.opened {
return nil
}
@ -197,22 +198,16 @@ func (db *DB) checkpoint() error {
var maxCheckpointedWALID int64
for {
// Determine last page of transaction.
metaWALID, metaFlags, err := db.findNextWALMetaPage(walID)
metaWALID, err := db.findNextWALMetaPage(walID)
if err == io.EOF {
break
} else if err != nil {
return err
}
// If transaction was rolled back, skip it.
if metaFlags&MetaPageFlagCommit == 0 {
walID = metaWALID + 1
continue
}
// Loop over pages in the transaction.
for ; walID <= metaWALID; walID++ {
canCheckpoint := minActiveWALID == 0 || walID <= minActiveWALID
canCheckpoint := exclusive || minActiveWALID == 0 || walID <= minActiveWALID
page, err := db.readWALPage(walID)
if err != nil {
@ -258,7 +253,7 @@ func (db *DB) checkpoint() error {
// Remove WAL segments that have been checkpointed.
if maxCheckpointedWALID != 0 {
for len(db.segments) > 1 {
for len(db.segments) > 0 {
segment := db.segments[0]
if segment.MaxWALID() >= maxCheckpointedWALID {
break
@ -274,21 +269,52 @@ func (db *DB) checkpoint() error {
}
}
// Ensure all segments are flushed and there is no remapped pages.
if exclusive {
assert(len(db.segments) == 0)
assert(pageMap.Len() == 0)
}
db.pageMap = pageMap
return nil
}
func (db *DB) findNextWALMetaPage(walID int64) (metaWALID int64, metaFlags uint32, err error) {
// truncateWALAfter removes all pages in the WAL after walID.
func (db *DB) truncateWALAfter(walID int64) error {
for i := len(db.segments) - 1; i >= 0; i-- {
segment := db.segments[i]
if segment.MaxWALID() <= walID {
break
}
// Drop entire segment if all pages are after WAL ID.
if walID < segment.MinWALID() {
if err := segment.Close(); err != nil {
return err
} else if err := os.Remove(segment.Path()); err != nil {
return err
}
db.segments, db.segments[i] = db.segments[:len(db.segments)-1], nil
continue
}
// If we only remove some of the WAL pages then truncate and exit
// since segments before this will retain all their pages.
return segment.TruncateAfter(walID)
}
return nil
}
func (db *DB) findNextWALMetaPage(walID int64) (metaWALID int64, err error) {
maxWALID := db.maxWALID()
for ; walID <= maxWALID; walID++ {
// Read page data from WAL and return if it is a meta page (either commit or rollback)
// Read page data from WAL and return if it is a meta page.
page, err := db.readWALPage(walID)
if err != nil {
return walID, metaFlags, err
return walID, err
} else if IsMetaPage(page) {
return walID, readFlags(page), nil
return walID, nil
}
// Skip over next page if this is a bitmap header.
@ -297,13 +323,30 @@ func (db *DB) findNextWALMetaPage(walID int64) (metaWALID int64, metaFlags uint3
}
}
return -1, 0, io.EOF
return -1, io.EOF
}
func (db *DB) findLastWALMetaPage() (walID int64, err error) {
if len(db.segments) == 0 {
return 0, nil
}
var maxWALID int64
for walID := db.segments[0].MinWALID(); walID <= maxWALID; walID++ {
if page, err := db.readWALPage(walID); err != nil {
return walID, err
} else if IsBitmapHeader(page) {
walID++ // skip next page for bitmap headers
} else if IsMetaPage(page) {
maxWALID = walID // save max meta WAL ID
}
}
return maxWALID, nil
}
// minActiveWALID returns the lowest WAL ID in use by any active transaction.
// Returns 0 if no transactions are active.
func (db *DB) minActiveWALID() int64 {
var walID int64
for tx := range db.txs {
if walID == 0 || walID > tx.walID {
@ -315,14 +358,12 @@ func (db *DB) minActiveWALID() int64 {
// ActiveWALSegment returns the most recent WAL segment.
func (db *DB) ActiveWALSegment() *WALSegment {
db.mu.RLock()
defer db.mu.RUnlock()
return db.activeWALSegment()
}
func (db *DB) activeWALSegment() *WALSegment {
if len(db.segments) == 0 {
return nil
}
@ -331,14 +372,12 @@ func (db *DB) activeWALSegment() *WALSegment {
// MinWALID returns the lowest WAL ID available in the WAL.
func (db *DB) MinWALID() int64 {
db.mu.RLock()
defer db.mu.RUnlock()
return db.minWALID()
}
func (db *DB) minWALID() int64 {
if len(db.segments) == 0 {
return 0
}
@ -347,7 +386,6 @@ func (db *DB) minWALID() int64 {
// MaxWALID returns the highest WAL ID available in the WAL.
func (db *DB) MaxWALID() int64 {
db.mu.RLock()
defer db.mu.RUnlock()
return db.maxWALID()
@ -364,7 +402,6 @@ func (db *DB) maxWALID() int64 {
// WALPageN returns the number of pages across all segments.
func (db *DB) WALPageN() int64 {
db.mu.RLock()
defer db.mu.RUnlock()
@ -385,8 +422,6 @@ func (db *DB) SyncWAL() error {
// readWALPage reads a single page at the given WAL ID.
func (db *DB) readWALPage(walID int64) ([]byte, error) {
//
// TODO(BBJ): Binary search for segment.
for _, s := range db.segments {
if walID >= s.MinWALID() && walID <= s.MaxWALID() {
@ -587,22 +622,71 @@ func (db *DB) initFreelistPage() error {
// Begin starts a new transaction.
func (db *DB) Begin(writable bool) (_ *Tx, err error) {
return db.begin(writable, false)
}
// TODO(BBJ): Acquire write lock if writable.
// BeginWithExclusiveLock starts a new transaction with an exclusive lock.
//
// This waits for all read transactions to finish and disallows any other
// transactions on the database. All WAL writes are flushed to disk and page
// writes during this transaction are written directly to the database file.
//
// Note that because page writes are direct, write failures can corrupt the
// database. This should only be used during bulk loading of data.
func (db *DB) BeginWithExclusiveLock() (_ *Tx, err error) {
return db.begin(true, true)
}
func (db *DB) begin(writable, exclusive bool) (_ *Tx, err error) {
if exclusive {
db.exclmu.Lock()
} else {
db.exclmu.RLock()
}
// Ensure only one writable transaction at a time.
if writable {
db.rwmu.Lock()
}
// This local function is called at exit points that occur before we can
// call Rollback() which would normally release these locks.
cleanup := func() {
if exclusive {
db.exclmu.Unlock()
} else {
db.exclmu.RUnlock()
}
if writable {
db.rwmu.Unlock()
}
}
db.mu.Lock()
defer db.mu.Unlock()
if !db.opened {
cleanup()
return nil, ErrClosed
}
tx := &Tx{db: db, rootRecords: db.rootRecords, pageMap: db.pageMap, writable: writable}
// Flush all WAL writes to disk before an exclusive writer so that we can
// work directly with the on-disk database.
if exclusive {
if err := db.checkpoint(true); err != nil {
cleanup()
return nil, err
}
}
tx := &Tx{
db: db,
rootRecords: db.rootRecords,
pageMap: db.pageMap,
writable: writable,
exclusive: exclusive,
}
// Copy meta page into transaction's buffer.
// This page is only written at the end of a dirty transaction.
@ -624,6 +708,12 @@ func (db *DB) Begin(writable bool) (_ *Tx, err error) {
// removeTx removes an active transaction from the database.
func (db *DB) removeTx(tx *Tx) error {
if tx.exclusive {
db.exclmu.Unlock()
} else {
db.exclmu.RUnlock()
}
// Release writer lock if tx is writable.
if tx.writable {
tx.db.rwmu.Unlock()
@ -635,7 +725,7 @@ func (db *DB) removeTx(tx *Tx) error {
// Write pages from WAL to DB.
// TODO(bbj): Move this to an async goroutine.
if tx.writable {
if err := db.checkpoint(); err != nil {
if err := db.checkpoint(false); err != nil {
return err
}
}

View file

@ -18,6 +18,7 @@ import (
"math/rand"
"os"
"testing"
"time"
"github.com/pilosa/pilosa/v2/rbf"
)
@ -129,3 +130,64 @@ func TestDB_Recovery(t *testing.T) {
tx.Rollback()
})
}
func TestDB_BeginWithExclusiveLock(t *testing.T) {
t.Run("EnsureBlock", func(t *testing.T) {
db := MustOpenDB(t)
defer MustCloseDB(t, db)
tx, err := db.BeginWithExclusiveLock()
if err != nil {
t.Fatal(err)
} else if err := tx.CreateBitmap("x"); err != nil {
t.Fatal(err)
}
// Attempt to start another transaction in a second goroutine.
ch := make(chan struct{})
go func() {
tx1, err := db.Begin(false)
if err != nil {
panic(err)
}
defer tx1.Rollback()
close(ch) // signal
}()
// Ensure other transctions are blocked during an exclusive lock.
select {
case <-ch:
t.Fatal("secondary transaction too soon")
case <-time.After(100 * time.Millisecond):
}
// Release exclusive lock.
if err := tx.Commit(); err != nil {
t.Fatal(err)
}
// Ensure other transaction to begin after exclusive lock released.
select {
case <-time.After(1 * time.Second):
t.Fatal("expected secondary transaction")
case <-ch:
}
})
t.Run("EnsureNoWAL", func(t *testing.T) {
db := MustOpenDB(t)
defer MustCloseDB(t, db)
tx, err := db.BeginWithExclusiveLock()
if err != nil {
t.Fatal(err)
}
defer tx.Rollback()
if err := tx.CreateBitmap("x"); err != nil {
t.Fatal(err)
} else if got, want := db.WALSize(), int64(0); got != want {
t.Fatalf("WALSize()=%d, want %d", got, want)
}
})
}

View file

@ -38,6 +38,7 @@ type Tx struct {
rootRecords []*RootRecord // read-only cache of root records
pageMap *immutable.Map // mapping of database pages to WAL IDs
writable bool // if true, tx can write
exclusive bool // if true, tx writes directly to db file (no wal)
dirty bool // if true, changes have been made
// If Rollback() has already completed, don't do it again.
@ -109,13 +110,9 @@ func (tx *Tx) Rollback() {
// TODO(bbj): Invalidate DB if rollback fails. Possibly attempt reopen?
// If any pages have been written, ensure we write a new meta page with
// the rollback flag to mark the end of the transaction. This allows us to
// discard pages in the transaction during playback of the WAL on open.
// Remove all WAL pages that have been written by this transaction.
if tx.dirty {
if err := tx.writeMetaPage(MetaPageFlagRollback); err != nil {
panic(err)
} else if err := tx.db.SyncWAL(); err != nil {
if err := tx.db.truncateWALAfter(tx.walID); err != nil {
panic(err)
}
}
@ -924,31 +921,40 @@ func (tx *Tx) readPage(pgno uint32) ([]byte, error) {
}
func (tx *Tx) writePage(page []byte) error {
// fmt.Println("writePage", readPageNo(page))
// Mark transaction as dirty so we write a meta page on commit/rollback.
tx.dirty = true
// If we are running in exclusive mode, directly write page to database.
if tx.exclusive {
return tx.db.writePage(readPageNo(page), page)
}
// Write page to WAL and obtain position in WAL.
walID, err := tx.db.writeWALPage(page, false)
if err != nil {
return err
}
// Mark transaction as dirty so we write a meta page on commit/rollback.
tx.dirty = true
// Update page map with WAL position.
tx.pageMap = tx.pageMap.Set(readPageNo(page), walID)
return nil
}
func (tx *Tx) writeBitmapPage(pgno uint32, page []byte) error {
// Mark transaction as dirty so we write a meta page on commit/rollback.
tx.dirty = true
// If we are running in exclusive mode, directly write page to database.
if tx.exclusive {
return tx.db.writePage(pgno, page)
}
// Write bitmap to WAL and obtain WAL position of the actual page data (not the prefix page).
walID, err := tx.db.writeBitmapPage(pgno, page)
if err != nil {
return err
}
// Mark transaction as dirty so we write a meta page on commit/rollback.
tx.dirty = true
// Update page map with WAL position.
tx.pageMap = tx.pageMap.Set(pgno, walID)
return nil
@ -958,6 +964,11 @@ func (tx *Tx) writeMetaPage(flag uint32) error {
// Set meta flags.
writeFlags(tx.meta[:], flag)
// If we are running in exclusive mode, directly write page to database.
if tx.exclusive {
return tx.db.writePage(0, tx.meta[:])
}
// Write page to WAL and obtain position in WAL.
walID, err := tx.db.writeWALPage(tx.meta[:], true)
if err != nil {

View file

@ -222,6 +222,34 @@ func (s *WALSegment) WriteWALPage(page []byte, isMeta bool) (walID int64, err er
return walID, nil
}
// TruncateAfter removes all pages after a given WAL ID.
func (s *WALSegment) TruncateAfter(walID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
// Ensure this is a partial truncation. Full truncation of a segment
// should be performed by the DB since it needs to remove the segment.
assert(walID > s.minWALID)
// Update to new page size.
newPageN := int((walID - s.minWALID) + 1) // new page count of segment
truncPageN := newPageN - s.pageN // number of pages removed
s.pageN = newPageN
// Check to see if we are only truncating from the write cache.
writeCachePageN := len(s.writeCache) / PageSize
if truncPageN <= int(writeCachePageN) {
s.writeCache = s.writeCache[:(writeCachePageN-truncPageN)*PageSize]
return nil
}
// Clear write cache.
s.writeCache = s.writeCache[:0]
// Remove on disk pages.
return os.Truncate(s.path, int64(s.pageN)*PageSize)
}
// Flush flushes the write buffer to the OS cache.
func (s *WALSegment) Flush() error {
s.mu.Lock()
@ -254,31 +282,6 @@ func (s *WALSegment) sync() error {
return s.w.Sync()
}
// trimBitmapHeaderTrailer removes the last page if the last page is a bitmap header.
// This should only be called on the last segment during recovery. A bitmap
// header write is a 2-page write so a partial write would corrupt the WAL.
func (s *WALSegment) trimBitmapHeaderTrailer() error {
// Skip if there are no pages in this segment.
if s.PageN() == 0 {
return nil
}
// Skip if this is not a bitmap header page.
if page, err := s.ReadWALPage(s.MaxWALID()); err != nil {
return err
} else if !IsBitmapHeader(page) {
return nil
}
// Truncate last page and reduce page count.
if err := os.Truncate(s.Path(), s.Size()-PageSize); err != nil {
return err
}
s.pageN--
return nil
}
// FormatWALSegmentPath returns a path for a WAL segment using a WAL ID.
func FormatWALSegmentPath(walID int64) string {
return fmt.Sprintf("%016x.wal", walID)