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.

397 lines
11KB

  1. package bbolt
  2. import (
  3. "bytes"
  4. "fmt"
  5. "sort"
  6. )
  7. // Cursor represents an iterator that can traverse over all key/value pairs in a bucket in sorted order.
  8. // Cursors see nested buckets with value == nil.
  9. // Cursors can be obtained from a transaction and are valid as long as the transaction is open.
  10. //
  11. // Keys and values returned from the cursor are only valid for the life of the transaction.
  12. //
  13. // Changing data while traversing with a cursor may cause it to be invalidated
  14. // and return unexpected keys and/or values. You must reposition your cursor
  15. // after mutating data.
  16. type Cursor struct {
  17. bucket *Bucket
  18. stack []elemRef
  19. }
  20. // Bucket returns the bucket that this cursor was created from.
  21. func (c *Cursor) Bucket() *Bucket {
  22. return c.bucket
  23. }
  24. // First moves the cursor to the first item in the bucket and returns its key and value.
  25. // If the bucket is empty then a nil key and value are returned.
  26. // The returned key and value are only valid for the life of the transaction.
  27. func (c *Cursor) First() (key []byte, value []byte) {
  28. _assert(c.bucket.tx.db != nil, "tx closed")
  29. c.stack = c.stack[:0]
  30. p, n := c.bucket.pageNode(c.bucket.root)
  31. c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
  32. c.first()
  33. // If we land on an empty page then move to the next value.
  34. // https://github.com/boltdb/bolt/issues/450
  35. if c.stack[len(c.stack)-1].count() == 0 {
  36. c.next()
  37. }
  38. k, v, flags := c.keyValue()
  39. if (flags & uint32(bucketLeafFlag)) != 0 {
  40. return k, nil
  41. }
  42. return k, v
  43. }
  44. // Last moves the cursor to the last item in the bucket and returns its key and value.
  45. // If the bucket is empty then a nil key and value are returned.
  46. // The returned key and value are only valid for the life of the transaction.
  47. func (c *Cursor) Last() (key []byte, value []byte) {
  48. _assert(c.bucket.tx.db != nil, "tx closed")
  49. c.stack = c.stack[:0]
  50. p, n := c.bucket.pageNode(c.bucket.root)
  51. ref := elemRef{page: p, node: n}
  52. ref.index = ref.count() - 1
  53. c.stack = append(c.stack, ref)
  54. c.last()
  55. k, v, flags := c.keyValue()
  56. if (flags & uint32(bucketLeafFlag)) != 0 {
  57. return k, nil
  58. }
  59. return k, v
  60. }
  61. // Next moves the cursor to the next item in the bucket and returns its key and value.
  62. // If the cursor is at the end of the bucket then a nil key and value are returned.
  63. // The returned key and value are only valid for the life of the transaction.
  64. func (c *Cursor) Next() (key []byte, value []byte) {
  65. _assert(c.bucket.tx.db != nil, "tx closed")
  66. k, v, flags := c.next()
  67. if (flags & uint32(bucketLeafFlag)) != 0 {
  68. return k, nil
  69. }
  70. return k, v
  71. }
  72. // Prev moves the cursor to the previous item in the bucket and returns its key and value.
  73. // If the cursor is at the beginning of the bucket then a nil key and value are returned.
  74. // The returned key and value are only valid for the life of the transaction.
  75. func (c *Cursor) Prev() (key []byte, value []byte) {
  76. _assert(c.bucket.tx.db != nil, "tx closed")
  77. // Attempt to move back one element until we're successful.
  78. // Move up the stack as we hit the beginning of each page in our stack.
  79. for i := len(c.stack) - 1; i >= 0; i-- {
  80. elem := &c.stack[i]
  81. if elem.index > 0 {
  82. elem.index--
  83. break
  84. }
  85. c.stack = c.stack[:i]
  86. }
  87. // If we've hit the end then return nil.
  88. if len(c.stack) == 0 {
  89. return nil, nil
  90. }
  91. // Move down the stack to find the last element of the last leaf under this branch.
  92. c.last()
  93. k, v, flags := c.keyValue()
  94. if (flags & uint32(bucketLeafFlag)) != 0 {
  95. return k, nil
  96. }
  97. return k, v
  98. }
  99. // Seek moves the cursor to a given key and returns it.
  100. // If the key does not exist then the next key is used. If no keys
  101. // follow, a nil key is returned.
  102. // The returned key and value are only valid for the life of the transaction.
  103. func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
  104. k, v, flags := c.seek(seek)
  105. // If we ended up after the last element of a page then move to the next one.
  106. if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
  107. k, v, flags = c.next()
  108. }
  109. if k == nil {
  110. return nil, nil
  111. } else if (flags & uint32(bucketLeafFlag)) != 0 {
  112. return k, nil
  113. }
  114. return k, v
  115. }
  116. // Delete removes the current key/value under the cursor from the bucket.
  117. // Delete fails if current key/value is a bucket or if the transaction is not writable.
  118. func (c *Cursor) Delete() error {
  119. if c.bucket.tx.db == nil {
  120. return ErrTxClosed
  121. } else if !c.bucket.Writable() {
  122. return ErrTxNotWritable
  123. }
  124. key, _, flags := c.keyValue()
  125. // Return an error if current value is a bucket.
  126. if (flags & bucketLeafFlag) != 0 {
  127. return ErrIncompatibleValue
  128. }
  129. c.node().del(key)
  130. return nil
  131. }
  132. // seek moves the cursor to a given key and returns it.
  133. // If the key does not exist then the next key is used.
  134. func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
  135. _assert(c.bucket.tx.db != nil, "tx closed")
  136. // Start from root page/node and traverse to correct page.
  137. c.stack = c.stack[:0]
  138. c.search(seek, c.bucket.root)
  139. // If this is a bucket then return a nil value.
  140. return c.keyValue()
  141. }
  142. // first moves the cursor to the first leaf element under the last page in the stack.
  143. func (c *Cursor) first() {
  144. for {
  145. // Exit when we hit a leaf page.
  146. var ref = &c.stack[len(c.stack)-1]
  147. if ref.isLeaf() {
  148. break
  149. }
  150. // Keep adding pages pointing to the first element to the stack.
  151. var pgid pgid
  152. if ref.node != nil {
  153. pgid = ref.node.inodes[ref.index].pgid
  154. } else {
  155. pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
  156. }
  157. p, n := c.bucket.pageNode(pgid)
  158. c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
  159. }
  160. }
  161. // last moves the cursor to the last leaf element under the last page in the stack.
  162. func (c *Cursor) last() {
  163. for {
  164. // Exit when we hit a leaf page.
  165. ref := &c.stack[len(c.stack)-1]
  166. if ref.isLeaf() {
  167. break
  168. }
  169. // Keep adding pages pointing to the last element in the stack.
  170. var pgid pgid
  171. if ref.node != nil {
  172. pgid = ref.node.inodes[ref.index].pgid
  173. } else {
  174. pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
  175. }
  176. p, n := c.bucket.pageNode(pgid)
  177. var nextRef = elemRef{page: p, node: n}
  178. nextRef.index = nextRef.count() - 1
  179. c.stack = append(c.stack, nextRef)
  180. }
  181. }
  182. // next moves to the next leaf element and returns the key and value.
  183. // If the cursor is at the last leaf element then it stays there and returns nil.
  184. func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
  185. for {
  186. // Attempt to move over one element until we're successful.
  187. // Move up the stack as we hit the end of each page in our stack.
  188. var i int
  189. for i = len(c.stack) - 1; i >= 0; i-- {
  190. elem := &c.stack[i]
  191. if elem.index < elem.count()-1 {
  192. elem.index++
  193. break
  194. }
  195. }
  196. // If we've hit the root page then stop and return. This will leave the
  197. // cursor on the last element of the last page.
  198. if i == -1 {
  199. return nil, nil, 0
  200. }
  201. // Otherwise start from where we left off in the stack and find the
  202. // first element of the first leaf page.
  203. c.stack = c.stack[:i+1]
  204. c.first()
  205. // If this is an empty page then restart and move back up the stack.
  206. // https://github.com/boltdb/bolt/issues/450
  207. if c.stack[len(c.stack)-1].count() == 0 {
  208. continue
  209. }
  210. return c.keyValue()
  211. }
  212. }
  213. // search recursively performs a binary search against a given page/node until it finds a given key.
  214. func (c *Cursor) search(key []byte, pgid pgid) {
  215. p, n := c.bucket.pageNode(pgid)
  216. if p != nil && (p.flags&(branchPageFlag|leafPageFlag)) == 0 {
  217. panic(fmt.Sprintf("invalid page type: %d: %x", p.id, p.flags))
  218. }
  219. e := elemRef{page: p, node: n}
  220. c.stack = append(c.stack, e)
  221. // If we're on a leaf page/node then find the specific node.
  222. if e.isLeaf() {
  223. c.nsearch(key)
  224. return
  225. }
  226. if n != nil {
  227. c.searchNode(key, n)
  228. return
  229. }
  230. c.searchPage(key, p)
  231. }
  232. func (c *Cursor) searchNode(key []byte, n *node) {
  233. var exact bool
  234. index := sort.Search(len(n.inodes), func(i int) bool {
  235. // TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
  236. // sort.Search() finds the lowest index where f() != -1 but we need the highest index.
  237. ret := bytes.Compare(n.inodes[i].key, key)
  238. if ret == 0 {
  239. exact = true
  240. }
  241. return ret != -1
  242. })
  243. if !exact && index > 0 {
  244. index--
  245. }
  246. c.stack[len(c.stack)-1].index = index
  247. // Recursively search to the next page.
  248. c.search(key, n.inodes[index].pgid)
  249. }
  250. func (c *Cursor) searchPage(key []byte, p *page) {
  251. // Binary search for the correct range.
  252. inodes := p.branchPageElements()
  253. var exact bool
  254. index := sort.Search(int(p.count), func(i int) bool {
  255. // TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
  256. // sort.Search() finds the lowest index where f() != -1 but we need the highest index.
  257. ret := bytes.Compare(inodes[i].key(), key)
  258. if ret == 0 {
  259. exact = true
  260. }
  261. return ret != -1
  262. })
  263. if !exact && index > 0 {
  264. index--
  265. }
  266. c.stack[len(c.stack)-1].index = index
  267. // Recursively search to the next page.
  268. c.search(key, inodes[index].pgid)
  269. }
  270. // nsearch searches the leaf node on the top of the stack for a key.
  271. func (c *Cursor) nsearch(key []byte) {
  272. e := &c.stack[len(c.stack)-1]
  273. p, n := e.page, e.node
  274. // If we have a node then search its inodes.
  275. if n != nil {
  276. index := sort.Search(len(n.inodes), func(i int) bool {
  277. return bytes.Compare(n.inodes[i].key, key) != -1
  278. })
  279. e.index = index
  280. return
  281. }
  282. // If we have a page then search its leaf elements.
  283. inodes := p.leafPageElements()
  284. index := sort.Search(int(p.count), func(i int) bool {
  285. return bytes.Compare(inodes[i].key(), key) != -1
  286. })
  287. e.index = index
  288. }
  289. // keyValue returns the key and value of the current leaf element.
  290. func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
  291. ref := &c.stack[len(c.stack)-1]
  292. // If the cursor is pointing to the end of page/node then return nil.
  293. if ref.count() == 0 || ref.index >= ref.count() {
  294. return nil, nil, 0
  295. }
  296. // Retrieve value from node.
  297. if ref.node != nil {
  298. inode := &ref.node.inodes[ref.index]
  299. return inode.key, inode.value, inode.flags
  300. }
  301. // Or retrieve value from page.
  302. elem := ref.page.leafPageElement(uint16(ref.index))
  303. return elem.key(), elem.value(), elem.flags
  304. }
  305. // node returns the node that the cursor is currently positioned on.
  306. func (c *Cursor) node() *node {
  307. _assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")
  308. // If the top of the stack is a leaf node then just return it.
  309. if ref := &c.stack[len(c.stack)-1]; ref.node != nil && ref.isLeaf() {
  310. return ref.node
  311. }
  312. // Start from root and traverse down the hierarchy.
  313. var n = c.stack[0].node
  314. if n == nil {
  315. n = c.bucket.node(c.stack[0].page.id, nil)
  316. }
  317. for _, ref := range c.stack[:len(c.stack)-1] {
  318. _assert(!n.isLeaf, "expected branch node")
  319. n = n.childAt(int(ref.index))
  320. }
  321. _assert(n.isLeaf, "expected leaf node")
  322. return n
  323. }
  324. // elemRef represents a reference to an element on a given page/node.
  325. type elemRef struct {
  326. page *page
  327. node *node
  328. index int
  329. }
  330. // isLeaf returns whether the ref is pointing at a leaf page/node.
  331. func (r *elemRef) isLeaf() bool {
  332. if r.node != nil {
  333. return r.node.isLeaf
  334. }
  335. return (r.page.flags & leafPageFlag) != 0
  336. }
  337. // count returns the number of inodes or page elements.
  338. func (r *elemRef) count() int {
  339. if r.node != nil {
  340. return len(r.node.inodes)
  341. }
  342. return int(r.page.count)
  343. }