a simple url shortening service, in the same vein as bit.ly and tinyurl.com, written in Go and using BoltDB as a backend
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

708 lines
19KB

  1. package bbolt
  2. import (
  3. "fmt"
  4. "io"
  5. "os"
  6. "sort"
  7. "strings"
  8. "time"
  9. "unsafe"
  10. )
  11. // txid represents the internal transaction identifier.
  12. type txid uint64
  13. // Tx represents a read-only or read/write transaction on the database.
  14. // Read-only transactions can be used for retrieving values for keys and creating cursors.
  15. // Read/write transactions can create and remove buckets and create and remove keys.
  16. //
  17. // IMPORTANT: You must commit or rollback transactions when you are done with
  18. // them. Pages can not be reclaimed by the writer until no more transactions
  19. // are using them. A long running read transaction can cause the database to
  20. // quickly grow.
  21. type Tx struct {
  22. writable bool
  23. managed bool
  24. db *DB
  25. meta *meta
  26. root Bucket
  27. pages map[pgid]*page
  28. stats TxStats
  29. commitHandlers []func()
  30. // WriteFlag specifies the flag for write-related methods like WriteTo().
  31. // Tx opens the database file with the specified flag to copy the data.
  32. //
  33. // By default, the flag is unset, which works well for mostly in-memory
  34. // workloads. For databases that are much larger than available RAM,
  35. // set the flag to syscall.O_DIRECT to avoid trashing the page cache.
  36. WriteFlag int
  37. }
  38. // init initializes the transaction.
  39. func (tx *Tx) init(db *DB) {
  40. tx.db = db
  41. tx.pages = nil
  42. // Copy the meta page since it can be changed by the writer.
  43. tx.meta = &meta{}
  44. db.meta().copy(tx.meta)
  45. // Copy over the root bucket.
  46. tx.root = newBucket(tx)
  47. tx.root.bucket = &bucket{}
  48. *tx.root.bucket = tx.meta.root
  49. // Increment the transaction id and add a page cache for writable transactions.
  50. if tx.writable {
  51. tx.pages = make(map[pgid]*page)
  52. tx.meta.txid += txid(1)
  53. }
  54. }
  55. // ID returns the transaction id.
  56. func (tx *Tx) ID() int {
  57. return int(tx.meta.txid)
  58. }
  59. // DB returns a reference to the database that created the transaction.
  60. func (tx *Tx) DB() *DB {
  61. return tx.db
  62. }
  63. // Size returns current database size in bytes as seen by this transaction.
  64. func (tx *Tx) Size() int64 {
  65. return int64(tx.meta.pgid) * int64(tx.db.pageSize)
  66. }
  67. // Writable returns whether the transaction can perform write operations.
  68. func (tx *Tx) Writable() bool {
  69. return tx.writable
  70. }
  71. // Cursor creates a cursor associated with the root bucket.
  72. // All items in the cursor will return a nil value because all root bucket keys point to buckets.
  73. // The cursor is only valid as long as the transaction is open.
  74. // Do not use a cursor after the transaction is closed.
  75. func (tx *Tx) Cursor() *Cursor {
  76. return tx.root.Cursor()
  77. }
  78. // Stats retrieves a copy of the current transaction statistics.
  79. func (tx *Tx) Stats() TxStats {
  80. return tx.stats
  81. }
  82. // Bucket retrieves a bucket by name.
  83. // Returns nil if the bucket does not exist.
  84. // The bucket instance is only valid for the lifetime of the transaction.
  85. func (tx *Tx) Bucket(name []byte) *Bucket {
  86. return tx.root.Bucket(name)
  87. }
  88. // CreateBucket creates a new bucket.
  89. // Returns an error if the bucket already exists, if the bucket name is blank, or if the bucket name is too long.
  90. // The bucket instance is only valid for the lifetime of the transaction.
  91. func (tx *Tx) CreateBucket(name []byte) (*Bucket, error) {
  92. return tx.root.CreateBucket(name)
  93. }
  94. // CreateBucketIfNotExists creates a new bucket if it doesn't already exist.
  95. // Returns an error if the bucket name is blank, or if the bucket name is too long.
  96. // The bucket instance is only valid for the lifetime of the transaction.
  97. func (tx *Tx) CreateBucketIfNotExists(name []byte) (*Bucket, error) {
  98. return tx.root.CreateBucketIfNotExists(name)
  99. }
  100. // DeleteBucket deletes a bucket.
  101. // Returns an error if the bucket cannot be found or if the key represents a non-bucket value.
  102. func (tx *Tx) DeleteBucket(name []byte) error {
  103. return tx.root.DeleteBucket(name)
  104. }
  105. // ForEach executes a function for each bucket in the root.
  106. // If the provided function returns an error then the iteration is stopped and
  107. // the error is returned to the caller.
  108. func (tx *Tx) ForEach(fn func(name []byte, b *Bucket) error) error {
  109. return tx.root.ForEach(func(k, v []byte) error {
  110. return fn(k, tx.root.Bucket(k))
  111. })
  112. }
  113. // OnCommit adds a handler function to be executed after the transaction successfully commits.
  114. func (tx *Tx) OnCommit(fn func()) {
  115. tx.commitHandlers = append(tx.commitHandlers, fn)
  116. }
  117. // Commit writes all changes to disk and updates the meta page.
  118. // Returns an error if a disk write error occurs, or if Commit is
  119. // called on a read-only transaction.
  120. func (tx *Tx) Commit() error {
  121. _assert(!tx.managed, "managed tx commit not allowed")
  122. if tx.db == nil {
  123. return ErrTxClosed
  124. } else if !tx.writable {
  125. return ErrTxNotWritable
  126. }
  127. // TODO(benbjohnson): Use vectorized I/O to write out dirty pages.
  128. // Rebalance nodes which have had deletions.
  129. var startTime = time.Now()
  130. tx.root.rebalance()
  131. if tx.stats.Rebalance > 0 {
  132. tx.stats.RebalanceTime += time.Since(startTime)
  133. }
  134. // spill data onto dirty pages.
  135. startTime = time.Now()
  136. if err := tx.root.spill(); err != nil {
  137. tx.rollback()
  138. return err
  139. }
  140. tx.stats.SpillTime += time.Since(startTime)
  141. // Free the old root bucket.
  142. tx.meta.root.root = tx.root.root
  143. // Free the old freelist because commit writes out a fresh freelist.
  144. if tx.meta.freelist != pgidNoFreelist {
  145. tx.db.freelist.free(tx.meta.txid, tx.db.page(tx.meta.freelist))
  146. }
  147. if !tx.db.NoFreelistSync {
  148. err := tx.commitFreelist()
  149. if err != nil {
  150. return err
  151. }
  152. } else {
  153. tx.meta.freelist = pgidNoFreelist
  154. }
  155. // Write dirty pages to disk.
  156. startTime = time.Now()
  157. if err := tx.write(); err != nil {
  158. tx.rollback()
  159. return err
  160. }
  161. // If strict mode is enabled then perform a consistency check.
  162. // Only the first consistency error is reported in the panic.
  163. if tx.db.StrictMode {
  164. ch := tx.Check()
  165. var errs []string
  166. for {
  167. err, ok := <-ch
  168. if !ok {
  169. break
  170. }
  171. errs = append(errs, err.Error())
  172. }
  173. if len(errs) > 0 {
  174. panic("check fail: " + strings.Join(errs, "\n"))
  175. }
  176. }
  177. // Write meta to disk.
  178. if err := tx.writeMeta(); err != nil {
  179. tx.rollback()
  180. return err
  181. }
  182. tx.stats.WriteTime += time.Since(startTime)
  183. // Finalize the transaction.
  184. tx.close()
  185. // Execute commit handlers now that the locks have been removed.
  186. for _, fn := range tx.commitHandlers {
  187. fn()
  188. }
  189. return nil
  190. }
  191. func (tx *Tx) commitFreelist() error {
  192. // Allocate new pages for the new free list. This will overestimate
  193. // the size of the freelist but not underestimate the size (which would be bad).
  194. opgid := tx.meta.pgid
  195. p, err := tx.allocate((tx.db.freelist.size() / tx.db.pageSize) + 1)
  196. if err != nil {
  197. tx.rollback()
  198. return err
  199. }
  200. if err := tx.db.freelist.write(p); err != nil {
  201. tx.rollback()
  202. return err
  203. }
  204. tx.meta.freelist = p.id
  205. // If the high water mark has moved up then attempt to grow the database.
  206. if tx.meta.pgid > opgid {
  207. if err := tx.db.grow(int(tx.meta.pgid+1) * tx.db.pageSize); err != nil {
  208. tx.rollback()
  209. return err
  210. }
  211. }
  212. return nil
  213. }
  214. // Rollback closes the transaction and ignores all previous updates. Read-only
  215. // transactions must be rolled back and not committed.
  216. func (tx *Tx) Rollback() error {
  217. _assert(!tx.managed, "managed tx rollback not allowed")
  218. if tx.db == nil {
  219. return ErrTxClosed
  220. }
  221. tx.rollback()
  222. return nil
  223. }
  224. func (tx *Tx) rollback() {
  225. if tx.db == nil {
  226. return
  227. }
  228. if tx.writable {
  229. tx.db.freelist.rollback(tx.meta.txid)
  230. tx.db.freelist.reload(tx.db.page(tx.db.meta().freelist))
  231. }
  232. tx.close()
  233. }
  234. func (tx *Tx) close() {
  235. if tx.db == nil {
  236. return
  237. }
  238. if tx.writable {
  239. // Grab freelist stats.
  240. var freelistFreeN = tx.db.freelist.free_count()
  241. var freelistPendingN = tx.db.freelist.pending_count()
  242. var freelistAlloc = tx.db.freelist.size()
  243. // Remove transaction ref & writer lock.
  244. tx.db.rwtx = nil
  245. tx.db.rwlock.Unlock()
  246. // Merge statistics.
  247. tx.db.statlock.Lock()
  248. tx.db.stats.FreePageN = freelistFreeN
  249. tx.db.stats.PendingPageN = freelistPendingN
  250. tx.db.stats.FreeAlloc = (freelistFreeN + freelistPendingN) * tx.db.pageSize
  251. tx.db.stats.FreelistInuse = freelistAlloc
  252. tx.db.stats.TxStats.add(&tx.stats)
  253. tx.db.statlock.Unlock()
  254. } else {
  255. tx.db.removeTx(tx)
  256. }
  257. // Clear all references.
  258. tx.db = nil
  259. tx.meta = nil
  260. tx.root = Bucket{tx: tx}
  261. tx.pages = nil
  262. }
  263. // Copy writes the entire database to a writer.
  264. // This function exists for backwards compatibility.
  265. //
  266. // Deprecated; Use WriteTo() instead.
  267. func (tx *Tx) Copy(w io.Writer) error {
  268. _, err := tx.WriteTo(w)
  269. return err
  270. }
  271. // WriteTo writes the entire database to a writer.
  272. // If err == nil then exactly tx.Size() bytes will be written into the writer.
  273. func (tx *Tx) WriteTo(w io.Writer) (n int64, err error) {
  274. // Attempt to open reader with WriteFlag
  275. f, err := os.OpenFile(tx.db.path, os.O_RDONLY|tx.WriteFlag, 0)
  276. if err != nil {
  277. return 0, err
  278. }
  279. defer func() {
  280. if cerr := f.Close(); err == nil {
  281. err = cerr
  282. }
  283. }()
  284. // Generate a meta page. We use the same page data for both meta pages.
  285. buf := make([]byte, tx.db.pageSize)
  286. page := (*page)(unsafe.Pointer(&buf[0]))
  287. page.flags = metaPageFlag
  288. *page.meta() = *tx.meta
  289. // Write meta 0.
  290. page.id = 0
  291. page.meta().checksum = page.meta().sum64()
  292. nn, err := w.Write(buf)
  293. n += int64(nn)
  294. if err != nil {
  295. return n, fmt.Errorf("meta 0 copy: %s", err)
  296. }
  297. // Write meta 1 with a lower transaction id.
  298. page.id = 1
  299. page.meta().txid -= 1
  300. page.meta().checksum = page.meta().sum64()
  301. nn, err = w.Write(buf)
  302. n += int64(nn)
  303. if err != nil {
  304. return n, fmt.Errorf("meta 1 copy: %s", err)
  305. }
  306. // Move past the meta pages in the file.
  307. if _, err := f.Seek(int64(tx.db.pageSize*2), io.SeekStart); err != nil {
  308. return n, fmt.Errorf("seek: %s", err)
  309. }
  310. // Copy data pages.
  311. wn, err := io.CopyN(w, f, tx.Size()-int64(tx.db.pageSize*2))
  312. n += wn
  313. if err != nil {
  314. return n, err
  315. }
  316. return n, nil
  317. }
  318. // CopyFile copies the entire database to file at the given path.
  319. // A reader transaction is maintained during the copy so it is safe to continue
  320. // using the database while a copy is in progress.
  321. func (tx *Tx) CopyFile(path string, mode os.FileMode) error {
  322. f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_TRUNC, mode)
  323. if err != nil {
  324. return err
  325. }
  326. err = tx.Copy(f)
  327. if err != nil {
  328. _ = f.Close()
  329. return err
  330. }
  331. return f.Close()
  332. }
  333. // Check performs several consistency checks on the database for this transaction.
  334. // An error is returned if any inconsistency is found.
  335. //
  336. // It can be safely run concurrently on a writable transaction. However, this
  337. // incurs a high cost for large databases and databases with a lot of subbuckets
  338. // because of caching. This overhead can be removed if running on a read-only
  339. // transaction, however, it is not safe to execute other writer transactions at
  340. // the same time.
  341. func (tx *Tx) Check() <-chan error {
  342. ch := make(chan error)
  343. go tx.check(ch)
  344. return ch
  345. }
  346. func (tx *Tx) check(ch chan error) {
  347. // Force loading free list if opened in ReadOnly mode.
  348. tx.db.loadFreelist()
  349. // Check if any pages are double freed.
  350. freed := make(map[pgid]bool)
  351. all := make([]pgid, tx.db.freelist.count())
  352. tx.db.freelist.copyall(all)
  353. for _, id := range all {
  354. if freed[id] {
  355. ch <- fmt.Errorf("page %d: already freed", id)
  356. }
  357. freed[id] = true
  358. }
  359. // Track every reachable page.
  360. reachable := make(map[pgid]*page)
  361. reachable[0] = tx.page(0) // meta0
  362. reachable[1] = tx.page(1) // meta1
  363. if tx.meta.freelist != pgidNoFreelist {
  364. for i := uint32(0); i <= tx.page(tx.meta.freelist).overflow; i++ {
  365. reachable[tx.meta.freelist+pgid(i)] = tx.page(tx.meta.freelist)
  366. }
  367. }
  368. // Recursively check buckets.
  369. tx.checkBucket(&tx.root, reachable, freed, ch)
  370. // Ensure all pages below high water mark are either reachable or freed.
  371. for i := pgid(0); i < tx.meta.pgid; i++ {
  372. _, isReachable := reachable[i]
  373. if !isReachable && !freed[i] {
  374. ch <- fmt.Errorf("page %d: unreachable unfreed", int(i))
  375. }
  376. }
  377. // Close the channel to signal completion.
  378. close(ch)
  379. }
  380. func (tx *Tx) checkBucket(b *Bucket, reachable map[pgid]*page, freed map[pgid]bool, ch chan error) {
  381. // Ignore inline buckets.
  382. if b.root == 0 {
  383. return
  384. }
  385. // Check every page used by this bucket.
  386. b.tx.forEachPage(b.root, 0, func(p *page, _ int) {
  387. if p.id > tx.meta.pgid {
  388. ch <- fmt.Errorf("page %d: out of bounds: %d", int(p.id), int(b.tx.meta.pgid))
  389. }
  390. // Ensure each page is only referenced once.
  391. for i := pgid(0); i <= pgid(p.overflow); i++ {
  392. var id = p.id + i
  393. if _, ok := reachable[id]; ok {
  394. ch <- fmt.Errorf("page %d: multiple references", int(id))
  395. }
  396. reachable[id] = p
  397. }
  398. // We should only encounter un-freed leaf and branch pages.
  399. if freed[p.id] {
  400. ch <- fmt.Errorf("page %d: reachable freed", int(p.id))
  401. } else if (p.flags&branchPageFlag) == 0 && (p.flags&leafPageFlag) == 0 {
  402. ch <- fmt.Errorf("page %d: invalid type: %s", int(p.id), p.typ())
  403. }
  404. })
  405. // Check each bucket within this bucket.
  406. _ = b.ForEach(func(k, v []byte) error {
  407. if child := b.Bucket(k); child != nil {
  408. tx.checkBucket(child, reachable, freed, ch)
  409. }
  410. return nil
  411. })
  412. }
  413. // allocate returns a contiguous block of memory starting at a given page.
  414. func (tx *Tx) allocate(count int) (*page, error) {
  415. p, err := tx.db.allocate(tx.meta.txid, count)
  416. if err != nil {
  417. return nil, err
  418. }
  419. // Save to our page cache.
  420. tx.pages[p.id] = p
  421. // Update statistics.
  422. tx.stats.PageCount += count
  423. tx.stats.PageAlloc += count * tx.db.pageSize
  424. return p, nil
  425. }
  426. // write writes any dirty pages to disk.
  427. func (tx *Tx) write() error {
  428. // Sort pages by id.
  429. pages := make(pages, 0, len(tx.pages))
  430. for _, p := range tx.pages {
  431. pages = append(pages, p)
  432. }
  433. // Clear out page cache early.
  434. tx.pages = make(map[pgid]*page)
  435. sort.Sort(pages)
  436. // Write pages to disk in order.
  437. for _, p := range pages {
  438. size := (int(p.overflow) + 1) * tx.db.pageSize
  439. offset := int64(p.id) * int64(tx.db.pageSize)
  440. // Write out page in "max allocation" sized chunks.
  441. ptr := (*[maxAllocSize]byte)(unsafe.Pointer(p))
  442. for {
  443. // Limit our write to our max allocation size.
  444. sz := size
  445. if sz > maxAllocSize-1 {
  446. sz = maxAllocSize - 1
  447. }
  448. // Write chunk to disk.
  449. buf := ptr[:sz]
  450. if _, err := tx.db.ops.writeAt(buf, offset); err != nil {
  451. return err
  452. }
  453. // Update statistics.
  454. tx.stats.Write++
  455. // Exit inner for loop if we've written all the chunks.
  456. size -= sz
  457. if size == 0 {
  458. break
  459. }
  460. // Otherwise move offset forward and move pointer to next chunk.
  461. offset += int64(sz)
  462. ptr = (*[maxAllocSize]byte)(unsafe.Pointer(&ptr[sz]))
  463. }
  464. }
  465. // Ignore file sync if flag is set on DB.
  466. if !tx.db.NoSync || IgnoreNoSync {
  467. if err := fdatasync(tx.db); err != nil {
  468. return err
  469. }
  470. }
  471. // Put small pages back to page pool.
  472. for _, p := range pages {
  473. // Ignore page sizes over 1 page.
  474. // These are allocated using make() instead of the page pool.
  475. if int(p.overflow) != 0 {
  476. continue
  477. }
  478. buf := (*[maxAllocSize]byte)(unsafe.Pointer(p))[:tx.db.pageSize]
  479. // See https://go.googlesource.com/go/+/f03c9202c43e0abb130669852082117ca50aa9b1
  480. for i := range buf {
  481. buf[i] = 0
  482. }
  483. tx.db.pagePool.Put(buf)
  484. }
  485. return nil
  486. }
  487. // writeMeta writes the meta to the disk.
  488. func (tx *Tx) writeMeta() error {
  489. // Create a temporary buffer for the meta page.
  490. buf := make([]byte, tx.db.pageSize)
  491. p := tx.db.pageInBuffer(buf, 0)
  492. tx.meta.write(p)
  493. // Write the meta page to file.
  494. if _, err := tx.db.ops.writeAt(buf, int64(p.id)*int64(tx.db.pageSize)); err != nil {
  495. return err
  496. }
  497. if !tx.db.NoSync || IgnoreNoSync {
  498. if err := fdatasync(tx.db); err != nil {
  499. return err
  500. }
  501. }
  502. // Update statistics.
  503. tx.stats.Write++
  504. return nil
  505. }
  506. // page returns a reference to the page with a given id.
  507. // If page has been written to then a temporary buffered page is returned.
  508. func (tx *Tx) page(id pgid) *page {
  509. // Check the dirty pages first.
  510. if tx.pages != nil {
  511. if p, ok := tx.pages[id]; ok {
  512. return p
  513. }
  514. }
  515. // Otherwise return directly from the mmap.
  516. return tx.db.page(id)
  517. }
  518. // forEachPage iterates over every page within a given page and executes a function.
  519. func (tx *Tx) forEachPage(pgid pgid, depth int, fn func(*page, int)) {
  520. p := tx.page(pgid)
  521. // Execute function.
  522. fn(p, depth)
  523. // Recursively loop over children.
  524. if (p.flags & branchPageFlag) != 0 {
  525. for i := 0; i < int(p.count); i++ {
  526. elem := p.branchPageElement(uint16(i))
  527. tx.forEachPage(elem.pgid, depth+1, fn)
  528. }
  529. }
  530. }
  531. // Page returns page information for a given page number.
  532. // This is only safe for concurrent use when used by a writable transaction.
  533. func (tx *Tx) Page(id int) (*PageInfo, error) {
  534. if tx.db == nil {
  535. return nil, ErrTxClosed
  536. } else if pgid(id) >= tx.meta.pgid {
  537. return nil, nil
  538. }
  539. // Build the page info.
  540. p := tx.db.page(pgid(id))
  541. info := &PageInfo{
  542. ID: id,
  543. Count: int(p.count),
  544. OverflowCount: int(p.overflow),
  545. }
  546. // Determine the type (or if it's free).
  547. if tx.db.freelist.freed(pgid(id)) {
  548. info.Type = "free"
  549. } else {
  550. info.Type = p.typ()
  551. }
  552. return info, nil
  553. }
  554. // TxStats represents statistics about the actions performed by the transaction.
  555. type TxStats struct {
  556. // Page statistics.
  557. PageCount int // number of page allocations
  558. PageAlloc int // total bytes allocated
  559. // Cursor statistics.
  560. CursorCount int // number of cursors created
  561. // Node statistics
  562. NodeCount int // number of node allocations
  563. NodeDeref int // number of node dereferences
  564. // Rebalance statistics.
  565. Rebalance int // number of node rebalances
  566. RebalanceTime time.Duration // total time spent rebalancing
  567. // Split/Spill statistics.
  568. Split int // number of nodes split
  569. Spill int // number of nodes spilled
  570. SpillTime time.Duration // total time spent spilling
  571. // Write statistics.
  572. Write int // number of writes performed
  573. WriteTime time.Duration // total time spent writing to disk
  574. }
  575. func (s *TxStats) add(other *TxStats) {
  576. s.PageCount += other.PageCount
  577. s.PageAlloc += other.PageAlloc
  578. s.CursorCount += other.CursorCount
  579. s.NodeCount += other.NodeCount
  580. s.NodeDeref += other.NodeDeref
  581. s.Rebalance += other.Rebalance
  582. s.RebalanceTime += other.RebalanceTime
  583. s.Split += other.Split
  584. s.Spill += other.Spill
  585. s.SpillTime += other.SpillTime
  586. s.Write += other.Write
  587. s.WriteTime += other.WriteTime
  588. }
  589. // Sub calculates and returns the difference between two sets of transaction stats.
  590. // This is useful when obtaining stats at two different points and time and
  591. // you need the performance counters that occurred within that time span.
  592. func (s *TxStats) Sub(other *TxStats) TxStats {
  593. var diff TxStats
  594. diff.PageCount = s.PageCount - other.PageCount
  595. diff.PageAlloc = s.PageAlloc - other.PageAlloc
  596. diff.CursorCount = s.CursorCount - other.CursorCount
  597. diff.NodeCount = s.NodeCount - other.NodeCount
  598. diff.NodeDeref = s.NodeDeref - other.NodeDeref
  599. diff.Rebalance = s.Rebalance - other.Rebalance
  600. diff.RebalanceTime = s.RebalanceTime - other.RebalanceTime
  601. diff.Split = s.Split - other.Split
  602. diff.Spill = s.Spill - other.Spill
  603. diff.SpillTime = s.SpillTime - other.SpillTime
  604. diff.Write = s.Write - other.Write
  605. diff.WriteTime = s.WriteTime - other.WriteTime
  606. return diff
  607. }