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.

605 lines
16KB

  1. package bbolt
  2. import (
  3. "bytes"
  4. "fmt"
  5. "sort"
  6. "unsafe"
  7. )
  8. // node represents an in-memory, deserialized page.
  9. type node struct {
  10. bucket *Bucket
  11. isLeaf bool
  12. unbalanced bool
  13. spilled bool
  14. key []byte
  15. pgid pgid
  16. parent *node
  17. children nodes
  18. inodes inodes
  19. }
  20. // root returns the top-level node this node is attached to.
  21. func (n *node) root() *node {
  22. if n.parent == nil {
  23. return n
  24. }
  25. return n.parent.root()
  26. }
  27. // minKeys returns the minimum number of inodes this node should have.
  28. func (n *node) minKeys() int {
  29. if n.isLeaf {
  30. return 1
  31. }
  32. return 2
  33. }
  34. // size returns the size of the node after serialization.
  35. func (n *node) size() int {
  36. sz, elsz := pageHeaderSize, n.pageElementSize()
  37. for i := 0; i < len(n.inodes); i++ {
  38. item := &n.inodes[i]
  39. sz += elsz + len(item.key) + len(item.value)
  40. }
  41. return sz
  42. }
  43. // sizeLessThan returns true if the node is less than a given size.
  44. // This is an optimization to avoid calculating a large node when we only need
  45. // to know if it fits inside a certain page size.
  46. func (n *node) sizeLessThan(v int) bool {
  47. sz, elsz := pageHeaderSize, n.pageElementSize()
  48. for i := 0; i < len(n.inodes); i++ {
  49. item := &n.inodes[i]
  50. sz += elsz + len(item.key) + len(item.value)
  51. if sz >= v {
  52. return false
  53. }
  54. }
  55. return true
  56. }
  57. // pageElementSize returns the size of each page element based on the type of node.
  58. func (n *node) pageElementSize() int {
  59. if n.isLeaf {
  60. return leafPageElementSize
  61. }
  62. return branchPageElementSize
  63. }
  64. // childAt returns the child node at a given index.
  65. func (n *node) childAt(index int) *node {
  66. if n.isLeaf {
  67. panic(fmt.Sprintf("invalid childAt(%d) on a leaf node", index))
  68. }
  69. return n.bucket.node(n.inodes[index].pgid, n)
  70. }
  71. // childIndex returns the index of a given child node.
  72. func (n *node) childIndex(child *node) int {
  73. index := sort.Search(len(n.inodes), func(i int) bool { return bytes.Compare(n.inodes[i].key, child.key) != -1 })
  74. return index
  75. }
  76. // numChildren returns the number of children.
  77. func (n *node) numChildren() int {
  78. return len(n.inodes)
  79. }
  80. // nextSibling returns the next node with the same parent.
  81. func (n *node) nextSibling() *node {
  82. if n.parent == nil {
  83. return nil
  84. }
  85. index := n.parent.childIndex(n)
  86. if index >= n.parent.numChildren()-1 {
  87. return nil
  88. }
  89. return n.parent.childAt(index + 1)
  90. }
  91. // prevSibling returns the previous node with the same parent.
  92. func (n *node) prevSibling() *node {
  93. if n.parent == nil {
  94. return nil
  95. }
  96. index := n.parent.childIndex(n)
  97. if index == 0 {
  98. return nil
  99. }
  100. return n.parent.childAt(index - 1)
  101. }
  102. // put inserts a key/value.
  103. func (n *node) put(oldKey, newKey, value []byte, pgid pgid, flags uint32) {
  104. if pgid >= n.bucket.tx.meta.pgid {
  105. panic(fmt.Sprintf("pgid (%d) above high water mark (%d)", pgid, n.bucket.tx.meta.pgid))
  106. } else if len(oldKey) <= 0 {
  107. panic("put: zero-length old key")
  108. } else if len(newKey) <= 0 {
  109. panic("put: zero-length new key")
  110. }
  111. // Find insertion index.
  112. index := sort.Search(len(n.inodes), func(i int) bool { return bytes.Compare(n.inodes[i].key, oldKey) != -1 })
  113. // Add capacity and shift nodes if we don't have an exact match and need to insert.
  114. exact := (len(n.inodes) > 0 && index < len(n.inodes) && bytes.Equal(n.inodes[index].key, oldKey))
  115. if !exact {
  116. n.inodes = append(n.inodes, inode{})
  117. copy(n.inodes[index+1:], n.inodes[index:])
  118. }
  119. inode := &n.inodes[index]
  120. inode.flags = flags
  121. inode.key = newKey
  122. inode.value = value
  123. inode.pgid = pgid
  124. _assert(len(inode.key) > 0, "put: zero-length inode key")
  125. }
  126. // del removes a key from the node.
  127. func (n *node) del(key []byte) {
  128. // Find index of key.
  129. index := sort.Search(len(n.inodes), func(i int) bool { return bytes.Compare(n.inodes[i].key, key) != -1 })
  130. // Exit if the key isn't found.
  131. if index >= len(n.inodes) || !bytes.Equal(n.inodes[index].key, key) {
  132. return
  133. }
  134. // Delete inode from the node.
  135. n.inodes = append(n.inodes[:index], n.inodes[index+1:]...)
  136. // Mark the node as needing rebalancing.
  137. n.unbalanced = true
  138. }
  139. // read initializes the node from a page.
  140. func (n *node) read(p *page) {
  141. n.pgid = p.id
  142. n.isLeaf = ((p.flags & leafPageFlag) != 0)
  143. n.inodes = make(inodes, int(p.count))
  144. for i := 0; i < int(p.count); i++ {
  145. inode := &n.inodes[i]
  146. if n.isLeaf {
  147. elem := p.leafPageElement(uint16(i))
  148. inode.flags = elem.flags
  149. inode.key = elem.key()
  150. inode.value = elem.value()
  151. } else {
  152. elem := p.branchPageElement(uint16(i))
  153. inode.pgid = elem.pgid
  154. inode.key = elem.key()
  155. }
  156. _assert(len(inode.key) > 0, "read: zero-length inode key")
  157. }
  158. // Save first key so we can find the node in the parent when we spill.
  159. if len(n.inodes) > 0 {
  160. n.key = n.inodes[0].key
  161. _assert(len(n.key) > 0, "read: zero-length node key")
  162. } else {
  163. n.key = nil
  164. }
  165. }
  166. // write writes the items onto one or more pages.
  167. func (n *node) write(p *page) {
  168. // Initialize page.
  169. if n.isLeaf {
  170. p.flags |= leafPageFlag
  171. } else {
  172. p.flags |= branchPageFlag
  173. }
  174. if len(n.inodes) >= 0xFFFF {
  175. panic(fmt.Sprintf("inode overflow: %d (pgid=%d)", len(n.inodes), p.id))
  176. }
  177. p.count = uint16(len(n.inodes))
  178. // Stop here if there are no items to write.
  179. if p.count == 0 {
  180. return
  181. }
  182. // Loop over each item and write it to the page.
  183. b := (*[maxAllocSize]byte)(unsafe.Pointer(&p.ptr))[n.pageElementSize()*len(n.inodes):]
  184. for i, item := range n.inodes {
  185. _assert(len(item.key) > 0, "write: zero-length inode key")
  186. // Write the page element.
  187. if n.isLeaf {
  188. elem := p.leafPageElement(uint16(i))
  189. elem.pos = uint32(uintptr(unsafe.Pointer(&b[0])) - uintptr(unsafe.Pointer(elem)))
  190. elem.flags = item.flags
  191. elem.ksize = uint32(len(item.key))
  192. elem.vsize = uint32(len(item.value))
  193. } else {
  194. elem := p.branchPageElement(uint16(i))
  195. elem.pos = uint32(uintptr(unsafe.Pointer(&b[0])) - uintptr(unsafe.Pointer(elem)))
  196. elem.ksize = uint32(len(item.key))
  197. elem.pgid = item.pgid
  198. _assert(elem.pgid != p.id, "write: circular dependency occurred")
  199. }
  200. // If the length of key+value is larger than the max allocation size
  201. // then we need to reallocate the byte array pointer.
  202. //
  203. // See: https://github.com/boltdb/bolt/pull/335
  204. klen, vlen := len(item.key), len(item.value)
  205. if len(b) < klen+vlen {
  206. b = (*[maxAllocSize]byte)(unsafe.Pointer(&b[0]))[:]
  207. }
  208. // Write data for the element to the end of the page.
  209. copy(b[0:], item.key)
  210. b = b[klen:]
  211. copy(b[0:], item.value)
  212. b = b[vlen:]
  213. }
  214. // DEBUG ONLY: n.dump()
  215. }
  216. // split breaks up a node into multiple smaller nodes, if appropriate.
  217. // This should only be called from the spill() function.
  218. func (n *node) split(pageSize int) []*node {
  219. var nodes []*node
  220. node := n
  221. for {
  222. // Split node into two.
  223. a, b := node.splitTwo(pageSize)
  224. nodes = append(nodes, a)
  225. // If we can't split then exit the loop.
  226. if b == nil {
  227. break
  228. }
  229. // Set node to b so it gets split on the next iteration.
  230. node = b
  231. }
  232. return nodes
  233. }
  234. // splitTwo breaks up a node into two smaller nodes, if appropriate.
  235. // This should only be called from the split() function.
  236. func (n *node) splitTwo(pageSize int) (*node, *node) {
  237. // Ignore the split if the page doesn't have at least enough nodes for
  238. // two pages or if the nodes can fit in a single page.
  239. if len(n.inodes) <= (minKeysPerPage*2) || n.sizeLessThan(pageSize) {
  240. return n, nil
  241. }
  242. // Determine the threshold before starting a new node.
  243. var fillPercent = n.bucket.FillPercent
  244. if fillPercent < minFillPercent {
  245. fillPercent = minFillPercent
  246. } else if fillPercent > maxFillPercent {
  247. fillPercent = maxFillPercent
  248. }
  249. threshold := int(float64(pageSize) * fillPercent)
  250. // Determine split position and sizes of the two pages.
  251. splitIndex, _ := n.splitIndex(threshold)
  252. // Split node into two separate nodes.
  253. // If there's no parent then we'll need to create one.
  254. if n.parent == nil {
  255. n.parent = &node{bucket: n.bucket, children: []*node{n}}
  256. }
  257. // Create a new node and add it to the parent.
  258. next := &node{bucket: n.bucket, isLeaf: n.isLeaf, parent: n.parent}
  259. n.parent.children = append(n.parent.children, next)
  260. // Split inodes across two nodes.
  261. next.inodes = n.inodes[splitIndex:]
  262. n.inodes = n.inodes[:splitIndex]
  263. // Update the statistics.
  264. n.bucket.tx.stats.Split++
  265. return n, next
  266. }
  267. // splitIndex finds the position where a page will fill a given threshold.
  268. // It returns the index as well as the size of the first page.
  269. // This is only be called from split().
  270. func (n *node) splitIndex(threshold int) (index, sz int) {
  271. sz = pageHeaderSize
  272. // Loop until we only have the minimum number of keys required for the second page.
  273. for i := 0; i < len(n.inodes)-minKeysPerPage; i++ {
  274. index = i
  275. inode := n.inodes[i]
  276. elsize := n.pageElementSize() + len(inode.key) + len(inode.value)
  277. // If we have at least the minimum number of keys and adding another
  278. // node would put us over the threshold then exit and return.
  279. if i >= minKeysPerPage && sz+elsize > threshold {
  280. break
  281. }
  282. // Add the element size to the total size.
  283. sz += elsize
  284. }
  285. return
  286. }
  287. // spill writes the nodes to dirty pages and splits nodes as it goes.
  288. // Returns an error if dirty pages cannot be allocated.
  289. func (n *node) spill() error {
  290. var tx = n.bucket.tx
  291. if n.spilled {
  292. return nil
  293. }
  294. // Spill child nodes first. Child nodes can materialize sibling nodes in
  295. // the case of split-merge so we cannot use a range loop. We have to check
  296. // the children size on every loop iteration.
  297. sort.Sort(n.children)
  298. for i := 0; i < len(n.children); i++ {
  299. if err := n.children[i].spill(); err != nil {
  300. return err
  301. }
  302. }
  303. // We no longer need the child list because it's only used for spill tracking.
  304. n.children = nil
  305. // Split nodes into appropriate sizes. The first node will always be n.
  306. var nodes = n.split(tx.db.pageSize)
  307. for _, node := range nodes {
  308. // Add node's page to the freelist if it's not new.
  309. if node.pgid > 0 {
  310. tx.db.freelist.free(tx.meta.txid, tx.page(node.pgid))
  311. node.pgid = 0
  312. }
  313. // Allocate contiguous space for the node.
  314. p, err := tx.allocate((node.size() + tx.db.pageSize - 1) / tx.db.pageSize)
  315. if err != nil {
  316. return err
  317. }
  318. // Write the node.
  319. if p.id >= tx.meta.pgid {
  320. panic(fmt.Sprintf("pgid (%d) above high water mark (%d)", p.id, tx.meta.pgid))
  321. }
  322. node.pgid = p.id
  323. node.write(p)
  324. node.spilled = true
  325. // Insert into parent inodes.
  326. if node.parent != nil {
  327. var key = node.key
  328. if key == nil {
  329. key = node.inodes[0].key
  330. }
  331. node.parent.put(key, node.inodes[0].key, nil, node.pgid, 0)
  332. node.key = node.inodes[0].key
  333. _assert(len(node.key) > 0, "spill: zero-length node key")
  334. }
  335. // Update the statistics.
  336. tx.stats.Spill++
  337. }
  338. // If the root node split and created a new root then we need to spill that
  339. // as well. We'll clear out the children to make sure it doesn't try to respill.
  340. if n.parent != nil && n.parent.pgid == 0 {
  341. n.children = nil
  342. return n.parent.spill()
  343. }
  344. return nil
  345. }
  346. // rebalance attempts to combine the node with sibling nodes if the node fill
  347. // size is below a threshold or if there are not enough keys.
  348. func (n *node) rebalance() {
  349. if !n.unbalanced {
  350. return
  351. }
  352. n.unbalanced = false
  353. // Update statistics.
  354. n.bucket.tx.stats.Rebalance++
  355. // Ignore if node is above threshold (25%) and has enough keys.
  356. var threshold = n.bucket.tx.db.pageSize / 4
  357. if n.size() > threshold && len(n.inodes) > n.minKeys() {
  358. return
  359. }
  360. // Root node has special handling.
  361. if n.parent == nil {
  362. // If root node is a branch and only has one node then collapse it.
  363. if !n.isLeaf && len(n.inodes) == 1 {
  364. // Move root's child up.
  365. child := n.bucket.node(n.inodes[0].pgid, n)
  366. n.isLeaf = child.isLeaf
  367. n.inodes = child.inodes[:]
  368. n.children = child.children
  369. // Reparent all child nodes being moved.
  370. for _, inode := range n.inodes {
  371. if child, ok := n.bucket.nodes[inode.pgid]; ok {
  372. child.parent = n
  373. }
  374. }
  375. // Remove old child.
  376. child.parent = nil
  377. delete(n.bucket.nodes, child.pgid)
  378. child.free()
  379. }
  380. return
  381. }
  382. // If node has no keys then just remove it.
  383. if n.numChildren() == 0 {
  384. n.parent.del(n.key)
  385. n.parent.removeChild(n)
  386. delete(n.bucket.nodes, n.pgid)
  387. n.free()
  388. n.parent.rebalance()
  389. return
  390. }
  391. _assert(n.parent.numChildren() > 1, "parent must have at least 2 children")
  392. // Destination node is right sibling if idx == 0, otherwise left sibling.
  393. var target *node
  394. var useNextSibling = (n.parent.childIndex(n) == 0)
  395. if useNextSibling {
  396. target = n.nextSibling()
  397. } else {
  398. target = n.prevSibling()
  399. }
  400. // If both this node and the target node are too small then merge them.
  401. if useNextSibling {
  402. // Reparent all child nodes being moved.
  403. for _, inode := range target.inodes {
  404. if child, ok := n.bucket.nodes[inode.pgid]; ok {
  405. child.parent.removeChild(child)
  406. child.parent = n
  407. child.parent.children = append(child.parent.children, child)
  408. }
  409. }
  410. // Copy over inodes from target and remove target.
  411. n.inodes = append(n.inodes, target.inodes...)
  412. n.parent.del(target.key)
  413. n.parent.removeChild(target)
  414. delete(n.bucket.nodes, target.pgid)
  415. target.free()
  416. } else {
  417. // Reparent all child nodes being moved.
  418. for _, inode := range n.inodes {
  419. if child, ok := n.bucket.nodes[inode.pgid]; ok {
  420. child.parent.removeChild(child)
  421. child.parent = target
  422. child.parent.children = append(child.parent.children, child)
  423. }
  424. }
  425. // Copy over inodes to target and remove node.
  426. target.inodes = append(target.inodes, n.inodes...)
  427. n.parent.del(n.key)
  428. n.parent.removeChild(n)
  429. delete(n.bucket.nodes, n.pgid)
  430. n.free()
  431. }
  432. // Either this node or the target node was deleted from the parent so rebalance it.
  433. n.parent.rebalance()
  434. }
  435. // removes a node from the list of in-memory children.
  436. // This does not affect the inodes.
  437. func (n *node) removeChild(target *node) {
  438. for i, child := range n.children {
  439. if child == target {
  440. n.children = append(n.children[:i], n.children[i+1:]...)
  441. return
  442. }
  443. }
  444. }
  445. // dereference causes the node to copy all its inode key/value references to heap memory.
  446. // This is required when the mmap is reallocated so inodes are not pointing to stale data.
  447. func (n *node) dereference() {
  448. if n.key != nil {
  449. key := make([]byte, len(n.key))
  450. copy(key, n.key)
  451. n.key = key
  452. _assert(n.pgid == 0 || len(n.key) > 0, "dereference: zero-length node key on existing node")
  453. }
  454. for i := range n.inodes {
  455. inode := &n.inodes[i]
  456. key := make([]byte, len(inode.key))
  457. copy(key, inode.key)
  458. inode.key = key
  459. _assert(len(inode.key) > 0, "dereference: zero-length inode key")
  460. value := make([]byte, len(inode.value))
  461. copy(value, inode.value)
  462. inode.value = value
  463. }
  464. // Recursively dereference children.
  465. for _, child := range n.children {
  466. child.dereference()
  467. }
  468. // Update statistics.
  469. n.bucket.tx.stats.NodeDeref++
  470. }
  471. // free adds the node's underlying page to the freelist.
  472. func (n *node) free() {
  473. if n.pgid != 0 {
  474. n.bucket.tx.db.freelist.free(n.bucket.tx.meta.txid, n.bucket.tx.page(n.pgid))
  475. n.pgid = 0
  476. }
  477. }
  478. // dump writes the contents of the node to STDERR for debugging purposes.
  479. /*
  480. func (n *node) dump() {
  481. // Write node header.
  482. var typ = "branch"
  483. if n.isLeaf {
  484. typ = "leaf"
  485. }
  486. warnf("[NODE %d {type=%s count=%d}]", n.pgid, typ, len(n.inodes))
  487. // Write out abbreviated version of each item.
  488. for _, item := range n.inodes {
  489. if n.isLeaf {
  490. if item.flags&bucketLeafFlag != 0 {
  491. bucket := (*bucket)(unsafe.Pointer(&item.value[0]))
  492. warnf("+L %08x -> (bucket root=%d)", trunc(item.key, 4), bucket.root)
  493. } else {
  494. warnf("+L %08x -> %08x", trunc(item.key, 4), trunc(item.value, 4))
  495. }
  496. } else {
  497. warnf("+B %08x -> pgid=%d", trunc(item.key, 4), item.pgid)
  498. }
  499. }
  500. warn("")
  501. }
  502. */
  503. type nodes []*node
  504. func (s nodes) Len() int { return len(s) }
  505. func (s nodes) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
  506. func (s nodes) Less(i, j int) bool { return bytes.Compare(s[i].inodes[0].key, s[j].inodes[0].key) == -1 }
  507. // inode represents an internal node inside of a node.
  508. // It can be used to point to elements in a page or point
  509. // to an element which hasn't been added to a page yet.
  510. type inode struct {
  511. flags uint32
  512. pgid pgid
  513. key []byte
  514. value []byte
  515. }
  516. type inodes []inode