diff --git a/rbf/db.go b/rbf/db.go index fcb5f86a9..a286370d4 100644 --- a/rbf/db.go +++ b/rbf/db.go @@ -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 } } diff --git a/rbf/db_test.go b/rbf/db_test.go index ef4e75ade..667fbf54d 100644 --- a/rbf/db_test.go +++ b/rbf/db_test.go @@ -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) + } + }) +} diff --git a/rbf/tx.go b/rbf/tx.go index b955d00a4..9b7e69111 100644 --- a/rbf/tx.go +++ b/rbf/tx.go @@ -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 { diff --git a/rbf/wal.go b/rbf/wal.go index fe7fcd40f..9a50fdde2 100644 --- a/rbf/wal.go +++ b/rbf/wal.go @@ -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)