diff --git a/mpt/DESIGN.md b/mpt/DESIGN.md new file mode 100644 index 0000000..3848038 --- /dev/null +++ b/mpt/DESIGN.md @@ -0,0 +1,733 @@ +# Merkle Patricia Tree Storage + +This package implements a Merkle Patricia Tree (MPT) stored on disk. + +A Merkle Patricia Tree (MPT) is a map that stores key-value pairs, where +each key and value is an opaque 256-bit value (typically a SHA256 hash). +Analogous to a [transparent log](https://research.swtch.com/tlog), +an MPT can cryptographically prove that a given key-value pair exists +(or that a key does not exist) in a given tree root. +By recording the sequence of tree roots in a transparent log, +a server can publish an record of the history of a key-value database, +in such a way that auditors can check that the database was correct +at all times, and clients can be sure the responses they received +came from the recorded database history. + +## Motivation and Performance Goals {#perf} + +[Certificate Transparency](https://certificate.transparency.dev/) (CT) +records a transparent log of all issued HTTPS certificates, +but it is expensive for a domain owner to +read the entire log looking for certificates for one or a few domains. +It would be useful to provide a server mapping from +domain to certificate lists, backed by an auditable MPT to verify +(after the fact) that the server has been behaving correctly. + +CT logs contain about 2 billion entries and are growing by about 75 entries/second. +[Let's Encrypt](https://letsencrypt.org/) currently issues 90-day certificates +but plans to start issuing +[short-lived, six-day certificates](https://letsencrypt.org/2024/12/11/eoy-letter-2024/), +which they estimate could result in 20X as many certificates. +The number of domains does not change, +so the overall MPT size would stay at around 2 billion entries, +but the update rate would grow to perhaps 75×20 = 1,500 updates/second. +It would be good to be able to handle 75×50 = 3,750 updates/second to provide extra headroom. +We'd like to handle this load comfortably on one reasonably configured server. + +Let's assume we have a single server with a +decent amount of memory and a fast disk. +For example, a Google Cloud +“m4-megamem-56” server with 56 vCPUs, 744 GiB of memory +costs $3,230/month (before sustained use or negotiated discounts). +Connecting two 1 TiB “balanced hyperdisks” each with +20,000 IOPS and 1200 MiB/s throughput adds $418/month. +([Pricing calculator](https://cloud.google.com/products/calculator?hl=en&dl=CjhDaVJsWVdZME1EVXpaQzFoWW1NMUxUUTFZMlF0T1RBNE15MDRNRFZtT0RFM09UWmhNR0VRQVE9PRAKGiRCQzY2RkQ0NC1CMEZGLTRFN0UtODJBMC0zNkM4NkI0QTQ5RjU)) +For people running servers in their own data centers +or behind their own couches, a +[Thelio Astra with 512 GB of memory and 2 4TB NVMe SSDs](https://system76.com/desktops/thelio-astra-a1.1-n1/configure) +can be had for $5,462. + +We will aim to be able to store, update, and serve an MPT of +2 billion entries with 4,000 updates/second easily on those servers. + +If we are to serve 4,000 updates/second, each update must cost less than 250µs. +And if our budget is 20,000 I/O operations/second, +each update must cost less than 5 I/O operations. +A tree of two billion entries has height at least 30: +we don't have enough I/O budget even to read nodes from disk during updates, +not to mention lookups. +The inevitable conclusion is that we must keep the entire tree in memory, +streaming updates to disk in batches. +So we will do exactly that. + +A hybrid approach is also possible, where each lookup or set requires O(1) disk I/O operations +in exchange for not keeping leaf data in memory. +This reduces the memory overhead from 112 bytes per entry to 48 bytes per entry. +After describing the fully in-memory version, we describe the hybrid version. + +## Merkle Patricia Tree Overview {#mpt} + +An MPT starts with the concept of a binary tree of depth 256, where the key-value +pairs are stored in the leaves at depth 256, and a lookup proceeds by walking +left or right according to each of the 256 key bits. The root node represents +the empty key prefix, its children represent key bit prefixes 0 and 1, their +children represent key bit prefixes 00, 01, 10, 11, and so on: at depth d, +the nodes represent key prefixes of d bits. +The original Key Transparency system at Google used exactly this data structure, +a [Merkle-hashed binary radix tree](https://github.com/google/keytransparency/blob/master/docs/overview.md). +Since then, the transparency community has realized that it works better + to apply Merkle hashing to a Patricia tree, +which adds three optimizations to the binary radix tree. + +First, the tree is “path-compressed,” by removing inner nodes with a single child: +a node that would have pointed at a single-child node is replaced by its child, +recursively. Every node is therefore either a leaf or an inner node with two children. +The path compression ensures that there are exactly _N_ inner nodes for a tree with _N_+1 leaf nodes. + +Second, unlike in a normal binary tree, an inner node stores only the bit position +that determines whether a lookup should proceed to the left or right child. +A lookup walks inner nodes down to some leaf, checking one bit at each step. +Only upon reaching the leaf does it do a full key comparison. +If it takes _O_(_K_) time to compare two keys, a normal binary tree would +take _O_(_K_ log _N_) time for a walk; this optimization +cuts the time to _O_(_K_ + log _N_). +Furthermore, inner nodes need not store associated keys, +cutting the number of stored keys by a factor of two. + +Third, nodes are “joined” by merging one inner node and one leaf node into +a single stored node. +(After joining _N_ inner nodes to _N_ leaf nodes, that leaves one “leaf-only” +node not paired to an inner node, +but the node is still stored using the joined representation.) +Whether a stored node represents an inner node or leaf node depends +on how it is reached +while walking the tree. +This trick is not essential, but it simplifies storage management +to have only one type of stored node. + +The path-compression optimization implies that an inner node for key prefix _p_ exists +if and only if the tree contains at least one key with prefix _p_0 and at least one key with prefix _p_1. +That is, the specific inner nodes that exist in a Patricia depend only on which +keys are present in the tree, not on their insertion order. +This implies that we can batch or otherwise reorder insertions of distinct keys +without affecting the final tree structure. + +For more about standard Patricia trees, see TODO REFERENCE. + +Cominbing the Merkle and Patricia pieces, a Merkle Patricia Tree provides the following operations: + + - Set(key, val): add a new key-value pair to the map. + - Snap(version): set the tree's version and return the tree root's key prefix and hash. + - Prove(key): return a proof of the result of looking up a given key in the current snapshot. A separate library function `Verify` verifies a proof and returns the lookup result (whether the key was found and, if so, its associated value). + - Sync(): flush recent changes to disk. + +The recursive hash of an MPT is defined as follows: + + - The hash of a leaf node is the hash of its key and value. + - The hash of an inner node is the hash of its bit position and its left and right children's hashes. + +A proof confirming that a key-value pair exists in an MPT with a given recursive hash +is the value followed by the bit position and sibling hash for every inner node along the path +back to the root. +The key-value pair can be hashed to obtain the hash of the leaf node, +and then the running hash can be hashed with the parent's bit position +and sibling hash to obtain the hash of the next inner node toward the root. +(Whether the running hash is the left or right child is determined by checking +the specified bit of the key.) +Recomputing the root's actual hash proves the lookup. + +A proof denying that a target key exists in an MPT is almost identical. +It consists of the “other key” whose leaf would be found by looking up the key in the tree, +followed by the proof that that other key is in the tree. +The verification checks that the other key's proof is valid and also that +the target key and other key agree at every relevant bit position. + +An in-memory MPT implementation is in [mem.go](mem.go). +It was useful to write and debug that version before adding the +complexity of attempting to store the tree on disk. +If future algorithmic bugs are found, it may still be helpful to debug them in +that version first. +It may also be useful read and understand that implementation +before proceeding to the disk implementation. + +## Encoding Details {#encoding} + +Any MPT implementation must define the exact encodings it uses. +The encodings used by this package are as follows. + +### Tree Hashes + +An empty tree is a special case that is otherwise independent +of the tree hash definition. In this implementation, +the hash of an empty tree is SHA256(_e_), the hash of the empty string (e3b0c442...7852b855). + +The hash of a leaf node is the hash of the concatenation of the key and value +(both fixed-size 32-byte sequences). + +The hash of an inner node at bit position _b_ with left and right child hashes _left_ and _right_ +is SHA256(_left_ || _right_ || _b_) where _left_ and _right_ are 32-byte values and _b_ is a one-byte value. + +### Proofs + +Proofs are variable length strings beginning with the 8-byte sequence `mptproof`. + +In Go the verifier's signature is: + + func Verify(snap Snapshot, key Key, proof Proof) (val Val, ok bool, err error) + +The verifier is given a tree hash (called a snapshot), a specific key, and a proof, +and it returns three results: (1) the value associated with the key, +if the proof proved the existence of the key in the tree, +(2) whether the proved result confirms or denies the existence of the key, +and (3) an error if the proof was invalid or did not match the tree hash. + +A proof of the empty tree is `mptproof` followed by a 0x00 byte. +It only applies when the tree hash is the empty tree hash, +and it disproves the existence of all possible keys. + +A proof confirming the existence of a key starts with `mptproof` followed by a 0x01 byte +and then a 32-byte value _v_. +The hash of the key's leaf node can be recomputed as _h_ = SHA256(_key_ || _v_). +If the tree is a single node, that is the entire proof: the verifier must check that +_h_ = _snap_. +If the tree contains more than one node, the proof continues with +one or more descriptions of sibling nodes along the path back to the tree root. +Each sibling node is encoded as 33 bytes: a one-byte bit position _b_ +followed by a 32-byte sibling hash _sib_. +The verifier must check whether the _b_'th bit of _key_ is 0 or 1 +and then update the running tree hash accordingly: + + - _h_ = SHA256(_h_ || _sib_ || _b_) if bit _b_ of _key_ is 0, or + - _h_ = SHA256(_sib_ || _h_ || _b_) if bit _b_ of _key_ is 1. + +Then, as before, the recomputed tree hash _h_ can be +compared against the actual tree hash. + +A proof denying the existence of a key starts with `mptproof` followed by a 0x02 byte +and then a 32-byte key _k_ and 32-byte value _v_, +describing a leaf node with hash _h_ = SHA256(_k_ || _v_). +If the tree is a single node, that is the entire proof: the verifier +must check that _h_ = _snap_ and that _k_ ≠ _key_. +Otherwise the proof format contains one or more siblings +encoded exactly as in the the existence proofs. +Verification also proceeds as in the existence proofs, +checking along the way that _k_ and _key_ agree on every bit _b_. +(Otherwise the proof would not describe the path taken +to walk through the tree in search of _key_.) +At the end, the verifier must check that _h_ = _snap_ and that _k_ ≠ _key_. + +The worst case length of an existence proof is 8+1+32+33*256 = 8489 bytes, +although random keys will never produce a path of length 256. + +The worst case length of a non-existence proof is 8+1+64+33*255 = 8488 bytes. +A non-existence proof can only have 255 siblings because otherwise the proof +would describe a key _k_ that agrees with _key_ at all 256 bit positions, +but then _k_ ≠ _key_ could not be true. +Again, random keys will never produce a path length of 256. + +Note: This encoding is considerably more compact than some others. +For example, the Rust akd crate's [MembershipProof](https://docs.rs/akd/0.12.0/akd/struct.MembershipProof.html) is: + + pub struct MembershipProof { + pub label: NodeLabel, + pub hash_val: AzksValue, + pub sibling_proofs: Vec, + } + +[NodeLabel](https://docs.rs/akd/0.12.0/akd/struct.NodeLabel.html) is a key plus a bit length, 32+4 = 36 bytes. \ +[AzksValue](https://docs.rs/akd/0.12.0/akd/struct.AzksValue.html) is 32 bytes. \ +[SiblingProof](https://docs.rs/akd/0.12.0/akd/struct.SiblingProof.html) is: + + pub struct SiblingProof { + pub label: NodeLabel, + pub siblings: [AzksElement; 1], + pub direction: Direction, + } + +[AzksElement](https://docs.rs/akd/0.12.0/akd/struct.AzksElement.html) is a NodeLabel and AzksValue, 64 bytes. \ +[Direction](https://docs.rs/akd/0.12.0/akd/enum.Direction.html) is a single byte. + +So SiblingProof is 36+64+1 = 101 bytes, and the overall worst case MembershipProof, if there are 256 siblings, +is 36+32+101*256 = 25,924 bytes. +The largest contributor to the difference is that the siblings include two node labels +when zero node labels suffice. +The result is a factor of three in the size of the proofs generated (and sent over the network). + +## Tree Algorithms + +There are two potentially important computations MPT hashes +that can be done without creating an explicit tree representation. + +### Whole Tree Hash + +The first computation is a whole-tree hash, +meaning to calculate a tree hash from a sorted list of key, value pairs +(sorted in key order). +Obviously the tree could be constructed and then hashed, +but the hash can be calculated more directly as follows. + +The algorithm maintains a stack _s_, +and we will denote the top element by _s_[−1], the one below it by _s_[−2], and the one below that by _s_[−3]. +The algorithm is: + + func treehash(list) -> (hash) + s = {} + for each k, v in list + s = reduce(push(s, leaf(k, v))) + return stackhash(s) + + func leaf(k, v) -> (node) + return {key: k, bits: 256, hash: SHA256(k || v)} + + func reduce(s) -> (stack) + while len(s) >= 3 and overlap(s[-3], s[-2]) > overlap(s[-2], s[-1]) + s = push(s[:-3], merge(s[-3], s[-2]), s[-1]) + return s + + func stackhash(s) -> (hash) + if len(s) == 0 + return SHA256() + while len(s) >= 2 + s = push(s[:-2], merge(s[-2], s[-1])) + return s[-1].hash + + func overlap(x, y) -> (bool) + return number of bits in shared prefix of x and y (at most min(x.bits, y.bits)) + + func merge(x, y) -> (node) + b = overlap(x, y) + return {key: x.key, bits: b, hash: SHA256(x.hash || y.hash || b)} + +Treehash maintains a stack corresponding to completed subtrees. +Each step pushes a new leaf node onto the stack +and then reduces the stack by the following observation. +When the stack has top values _x_, _y_, _z_, +and _x_ and _y_ have more bits in common (join together deeper in the tree) than _y_ and _z_, +then _x_ and _y_ can be merged into a single subtree, +since _z_ and all keys that follow will only be joined to it higher up in the tree. +Once the entire key-value list has been pushed on to the stack and reduced, +all that remains is for stackhash to merge the final fringe up the right side into a complete tree. + +The run time of this algorithm is _N_ pushes of leaves, _N_ calls to reduce, +and _N_−1 total merges. +Each call to reduce ends in a failed overlap comparison, +and the number of successful overlap comparisons is less than _N_ +(since there are only _N_−1 merges), +so there are at most 2_N_ overlap comparisons, or 4_N_ overlap calls. +That's _O_(_N_ _K_) time, where _K_ is again the time for a key comparison. + +### Predicted Tree Hash + +The second computation is a predicted tree hash, +meaning to calculate a tree hash that would result +from starting with an existing tree and inserting a +sorted list of key, value pairs that may replace existing +nodes or add new ones. +The trick is that we want to compute this hash without editing the existing tree. + +If every node in the tree stored the key corresponding to that node, +then we could use the following algorithm to produce a sorted list of nodes +corresponding to subtrees of the existing tree or new leaves, +and then we could apply a variant of treehash to that list of nodes +to compute the overall tree hash. + + func updatehash(tree, list) -> (hash) + s, list = update({}, tree, list) + for k, v in list + s = reduce(push(s, leaf(k, v))) + return stackhash(s) + + func update(s, node, list) -> (stack, list) + // push modifications before node + while len(list) > 0 or list[0].key < node.key comparing only node.bits bits + k, v = list[0] + list = list[1:] + s = reduce(push(s, leaf(k, v))) + // push entire subtree if no modifications inside it + if len(list) == 0 or node.key < list[0].key comparing only node.bits bits + return reduce(push(s, node)), list + // replace leaf if node is a leaf + if node.bits == 256 + k, v = list[0] + list = list[1:] + return reduce(push(s, leaf(k, v))), list + // apply modifications within subtree + s, list = update(s, node.left, list) + s, list = update(s, node.right, list) + return s, list + +The only problem is that the second Patricia optimization +removed the .key field in the tree nodes. +The solution is that the third Patricia optimization added it back implicitly. +Each new combined node is constructed during the +insertion of a key _k_ to hold the leaf _k_, v_ +as well as being the new inner node +representing a previously unexamined bit _b_ that +separates _k_ from an existing subtree. +Both _k_ and the subtree have the same key prefix of _b_ bits, +so the inner node's key is simply _k_ truncated to _b_ bits. +So the key stored in the combined node turns out to +describe _both_ the inner node and the leaf. +We still don't want to use it as an inner node key +during lookups, as that would turn our _O_(_K_ + log _N_) +back into _O_(_K_ log _N_), where _K_ is a key comparison. +But it's there, and we can use it while computing the predicted tree hash. + +Let's say there are _N_ items in the list and _T_ items in the existing tree +and that the tree has height _O_(log _T_). +Then as a worst case we can estimate that each item requires enumerating _O_(log _T_) +subtree nodes to make room for the insertion of the new leaf, +a total of _O_(_N_ log _T_) nodes pushed onto the stack, +which will require _O_(_N_ _K_ log _T_) time for the comparisons in the reduction to a single tree hash. +It also required one comparison per update call, +and there were asymptotically the same number of update calls +as nodes pushed, so the total time is still _O_(_N_ _K_ log _T_). + +There may be some way to use the Patricia property to avoid +the repeated comparisons in the algorithm described above, +but it would not change the stack reduction time, so the +overall asymptotic runtime would remain. + +## Storage Overview {#storage} + +Although the working tree is stored in memory, +we of course want to recover from crashes by persisting the tree to a disk file as well. +The disk file consists of a memory image of a tree +followed by a sequence of patches of the form “at offset O, write these N bytes”. +Each update requires only a single disk write to append its patches to the file. +(Multiple updates can also be batched into a single write.) + +There is only one problem with this representation: it grows without bound, +and faster than the tree. +A write of an existing key-value pair requires no additional memory at all, +but it requires 30 or so patches to existing nodes, +to update the hashes of the inner nodes along the path back to the root. +A write of a new key-value pair requires only one new allocated node, +but it too requires the same 30 or so patches. +There must be some kind of compaction. + +The simplest way to compact one file is to write out a second disk file. +After all, the in-memory copy has all the patches applied already. +Conceptually, we can stop updates, +write the current tree memory to a new file, +delete the old file, and then resume updates, +now writing patches to the new file. +It is worth introducing two complications. +First, we can reuse the old file as the output for the next compaction, +alternating between a pair of files +instead of continually deleting and recreating files. +Second, we can let updates proceed concurrently +with compaction, so that updates aren't blocked +waiting to write a few hundred gigabytes to disk. + +## Memory Format {#mem} + +The memory format of the tree must be suitable for writing to disk +and then reading back into a different memory location, +so it cannot contain actual Go pointers. +Instead, the memory format is one very large byte array +that is interpreted as higher-level data structures +“on demand,” when accessing or modifying it. +Each actual update happens by writing to the memory +as bytes and also logging the mutation to a patch +that will be written to disk. + +The tree memory starts with a header with the form: + + version [ 8 bytes] + dirty [ 1 byte] + pad [ 1 byte] + root [ 6 bytes] + hash [32 bytes] + nodes [ 8 bytes] + +All numbers are stored in big-endian order +for legibility when reading hex dumps. + + - “version” is a number for clients to use to match the + tree contents to a position in the underlying transparent log. + - “root” is a pointer to the tree's root node, + represented as a 48-bit byte offset within the + tree memory. + - “nodes” field counts the number of nodes (leaves) + stored in the tree. + - “hash” is the Merkle hash of the tree root. + When “dirty” is set, the hash is stale and needs to be recomputed. + - “pad” pads “root” to a 16-bit boundary and “hash” and “nodes” + to a 64-bit boundary. + +The header is immediately followed by a sequence of Patricia nodes, +each with the form: + + key [32 bytes] + val [32 bytes] + bit [ 1 byte] + dirty [ 1 byte] + pad [ 2 bytes] + left [ 6 bytes] + right [ 6 bytes] + ihash [32 bytes] + +Remember that each Patricia node represents both one leaf node +and one inner node. + + - “key” and “val” are the key and value for the leaf node. + - “left” and “right” are pointers to the inner node's + left and right children; + “bit” is the bit position to use to decide between them + during a lookup. + Since keys are 32 bytes, bit positions 0..255 fit in a single byte. + The leaf-only node is an exception: it needs a bit position + set to -1, but it can be identified by having “left” and “right” set to 0 + (nil pointer) + and treated as a special case when reading “bit”. + - “ihash” is the Merkle hash of the inner node. + When “dirty” is set, the hash is stale and needs to be recomputed. + The Merkle hash of the leaf node is not stored explicitly. + It is recomputed from “key” and “val” whenever it is needed. + +After setting the value associated with a given key, +the “ihash” values in all nodes back to the root +need to be recomputed. +If we are writing _N_ new values between +taking snapshots, +we would end up recomputing the root hash _N_−1 times +unnecessarily, +recomputing the root's two children's hashes _N_/2−1 times +unnecessarily, and so on. +Since everything else is in memory, batched, and cheap, +these SHA256 computations end up being the serving bottleneck. +To avoid the unnecessary hashes, +we don't recompute any hash during a write of a new value. +Instead, we set the “dirty” field on all nodes back +to the root. +The snapshot operation restores each hash by recomputing it +from its children, +restoring those hashes first as needed. +In effect, it rewalks the entire modified area of the tree, +computing all the new hashes then. +Snapshots are still amortized O(1) but not an actual O(1). +If the snapshot operations caused problematic latency hiccups, +this lazy recomputation could be abandoned. + +Notice that a Patricia node takes 112 bytes, +so a 2-billion node tree requires about 224 GB of memory, +well within the 512 GB we allotted ourselves on our +“reasonably configured server”. +The actual memory for the tree is obtained directly +using the operating system, not from the Go heap. +Using _mmap_(2), we can reserve a very large amount +of space but then only map the memory we need as the tree grows. +This allows extending the tree without having to move it. +It also has the side benefit of not skewing the Go +garbage collector's pacing with one extremely large allocation. + +The file format allows grouping memory updates into atomic units, +so that partially applied updates are never observed +when loading a tree from disk. + +## File Format {#format} + +A file consists of the magic string `"mpt tree\n\x00\x00\x00\x00\x00\x00\x00"` +followed by a sequence of variable-length frames. +Each frame has the form: + + treeID [16 bytes] + treeSeq [ 8 bytes] + N [ 8 bytes] + data [ N bytes] + checksum [32 bytes] + +The “treeID” is randomly chosen when a tree is first created, +and the “treeSeq” is a sequence number incremented each time +a new tree file is written. +Both ensure that when files are reused +(either for a new tree or a new version of the same tree), +old frames are not misinterpreted as new ones. + +The “checksum” is a SHA256 checksum of the preceding fields. +Verifying the checksum detects corruption but also provides +atomicity of frame writes: either the whole frame is written +to disk and the checksum matches, or none of it is used. + +The file starts with one very large frame containing a snapshot +of the tree memory. +As we will see, concurrent compaction means that +this memory snapshot may not itself be a valid tree: +the patches in the rest of the file must be applied +not just to obtain the latest tree, but also to obtain +a valid tree. + +The second and subsequent frames in the file each hold +a patch block, which holds one or more mutations of the form: + + offset [varint] + N [varint] + data [N bytes] + +That mutation says to write `data` of length `N` at +`offset` in the tree memory. + +There is no guarantee that a file ends after a valid patch block frame. +If a frame was only partially written before a process or system crash, +we still want to read the tree before that point. +We do this by reading as many valid (checksum-matching) frames +as possible from the file and stopping at EOF or when we reach +a frame that is truncated or does not have a valid checksum. + +## Compaction {#compaction} + +When the current disk file holding a tree has grown too large, +compaction writes and then switches to a new smaller file, +at which point the old one can be abandoned. +In practice, the implementation reuses the old one for the +next compaction. + +Logically, compaction is shortening the old file by writing +the patches directly to the tree memory. +However, it is not reading the old file: the tree +can be written directly from memory instead of consulting +the old file. +So technically compaction may be a misnomer: +the old file not being compacted so much as it is +being obsoleted. + +One approach would be to pause all tree updates, +write the tree to the new file, and then continue +updates, writing patches to the new file. +If the tree is 224 GB, +then even if we can write at a relatively fast 10 GB/s, +that would be a 22-second pause. +Instead, we can allow tree updates to proceed +concurrently with compaction. + +The concurrent compaction algorithm works as follows: + + - Record the current tree memory size _M_. + - Arrange to write future patch frames + to both the _current_ tree file and the _next_ tree file. + In the _next_ tree, the patches start at the offset where + a tree of size _M_ would end. + - Each time a new megabyte (or other chosen chunk size) of patches + is written, write the next megabyte of tree memory to the + _next_ tree file as well. + - Once all the tree memory has been written to the _next_ file, + sync it to disk and write its incremented tree sequence number + so that a future open will choose that file. + - Now the old _next_ file has become the new _current_ one, and the + old _current_ one will be reused for as _next_ for the next + compaction. + +Note that the compaction is concurrent but not parallel: +the compaction writes are interleaved with non-compaction writes, +not run in a separate goroutine. + +Suppose we start a compaction when the in-memory tree is _M_ bytes +and the _current_ tree file is 2*M* bytes long. +Compaction will finish when another _M_ bytes of patches have been written, +meaning the _current_ file will be 3*M* bytes when it is retired, +and the _next_ tree file will be 2*M* bytes long. +After installing _next_ as the new _current_, it will be time for +a new compaction. +The result is continuous, just-in-time compaction: each one finishes +exactly as the next needs to begin. + +For the actual implementation, +it seemed safer to write two megabytes of tree for +each megabyte of patches, rather than dance on that knife's edge. +If writing a megabyte of patches to two different trees +(_current_ and _next_) allows writing two megabytes of tree memory as well, +then the implementation writes the same amount of tree bytes and patch bytes to disk +when compaction is running. +At worst, compaction is always running, so that the number of tree +bytes written equals the number of patch bytes written. + +If each memory change is written to two patch files, +and those writes justify writing the same amount of tree bytes, +then the [write amplification](https://en.wikipedia.org/wiki/Write_amplification) +factor is no worse than 4. +This constant factor is much better than the _O_(log _N_) amplification in +[log-structured merge trees](https://en.wikipedia.org/wiki/Log-structured_merge-tree). +The improvement is possible because we keep all the data in memory at all times. + +## Speed {#speed} + +On my circa-2023 home server with 128 GB of RAM +and an NVMe disk using LVM encryption, storing 834 million hashes +takes about 250 minutes, or about 55,000 Set operations per second. +This is with constant disk compaction, and I suspect something in my +kernel stack of slowing disk I/O. + +A “lazy hash” optimization that delays recomputing all inner node hashes +is delayed until the Sync operation can avoid spending time +computing hashes that will be overwritten by a subsequent Set, +but it dramatically increases the latency of Sync. +More important than not computing the hashes is not writing +them to disk, especially for the somewhat special case of writing all new entries +when populating a new tree. +In that case, the lazy hash's lower disk usage +also avoids any compaction: only new data is +being written, so the disk file never reaches twice the memory size. +In that case, the 834 million hashes can be written in 45 minutes, +followed by a 7 minute sync, or about 260,000 Set operations second. + +A limited lazy hash that is lazy only up to a fixed number of +Set operations may be the best of both worlds. + +Prove operations run in microseconds. + +Snap is effectively free. + +## Recovery {#recovery} + +On crash and restart, the persistent disk implementation guarantees +to have a state equivalent to some prefix of the operations that had +been executed prior to the crash. That is, some recent operations may +have been lost, but if operation K is observed, then all operations +prior to K will also be observed. + +If the client calls Snap with a version number at regular intervals, +then after a crash, if the tree loads with version V, it means that +all of the Set calls before Snap(V) has been retained, and some of the Set +calls between Snap(V) and Snap(V+1) may also have been retained. +It suffices to replay all the Set operations between Snap(V) and Snap(V+1) +and then Snap(V+1) to get a consistent tree. + +## Hybrid Approach + +The approach described so far is the original in-memory approach. +It is tagged as mpt v0.1.0. + +This section describes a hybrid approach implemented in later versions. +The hybrid approach trades a constant number of disk I/O per Set or Prove +operation for reduced memory requirements. In the hybrid approach, the +leaf nodes (meaning the key and value fields) are all stored in a “leaf file” +not stored in memory. Writes to the leaf file are still recorded in patch blocks, +so that after recovery the leaf file is always at least as up to date as the main +tree memory image. However, writes to the leaf file also happen immediately, +so after recovery, the leaf file may also contain writes beyond those reflected +in the main tree memory image. Having a leaf file that is “too new” cannot +affect the structure of the overall tree, since keys are never changed after a +node is allocated. However, the leaf file being too new can mean that values +that are “too new” are recorded for leaf nodes, so the client must recover by +replaying all the Set operations that may have happened after the point +where the memory image was recovered. Once those are replayed, the +memory image and the leaf file will be in sync. + +To support the recovery operation, there is a new method Tree.Version: + + // Version returns the version number of the tree's last complete snapshot. + // All Set calls made prior to Snap(version) are guaranteed to be + // recorded in the tree. However, if exact is false, then the tree may + // include the effect of Set calls made after that snapshot. + // In that case, to bring the tree into a consistent state, the client is + // expected to replay all Set calls up to the next version. + Version() (version int64, exact bool) + +In this new approach, calling Prove requires two disk I/Os: one to read the +leaf key and value at the end of the lookup, and one to read that node's sibling +for inclusion in the proof. Calling set requires three disk I/Os: the same two reads +needed by Prove as well as one write to update or create a leaf. + +In exchange for these two or three disk I/Os per operation, the memory +requirements are reduced from 112 bytes per record to 48 bytes per record. diff --git a/mpt/disk.go b/mpt/disk.go new file mode 100644 index 0000000..6cb05b4 --- /dev/null +++ b/mpt/disk.go @@ -0,0 +1,504 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package mpt + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "os" + "runtime" + "sync" + + "filippo.io/torchwood/mpt/internal/pmem" +) + +// Tree Format +// +// The tree memory starts with a header: +// +// version [8] +// dirty [1] +// pad [1] +// root [6] +// hash [32] +// nodes [8] +// +// The header is followed by a sequence of Patricia nodes of the form: +// +// key [32] +// val [32] +// bit [1] +// dirty [1] +// pad [2] +// left [6] +// right [6] +// ihash [32] +// +// The root, left, and right “pointers” are byte offsets from the start of the tree memory. +// A nil pointer is stored as offset 0, which would otherwise point at the tree header. + +const ( + // header offsets + hdrVersion = 0 + hdrDirty = 8 + hdrExact = 9 + hdrRoot = 10 + hdrHash = 16 + hdrSize = 48 + + // node offsets + // setLeftRight knows that left and right are contiguous. + nodeUbit = 0 + nodeDirty = 1 + nodeLeft = 4 + nodeRight = 10 + nodeIHash = 16 + nodeSize = 48 + + // address size + addrSize = 6 + + // leaf file details + leafMagic = "mpt leaf\n\x00\x00\x00\x00\x00\x00\x00" + leafHdrSize = 16 + leafSize = 64 +) + +// File is the interface needed for on-disk storage. +type File interface { + io.ReaderAt + io.WriterAt + io.Closer + Sync() error +} + +// File must implement pmem.File. +// It really should be exactly pmem.File but we don't want to +// expose pmem in the API definitions, so File is a copy instead. +var _ pmem.File = File(nil) + +// A diskTree is an on-disk [Tree]. +type diskTree struct { + // mmu is the memory mapping mutex. + // All methods except Close do mmu.RLock and mmu.RUnlock + // in order to be allowed to use pmem.Data() aka mem. + // Close calls mmu.Lock/mmu.Unlock to wait for all other + // method calls to finish before unmapping the memory. + mmu sync.RWMutex + pmem *pmem.Mem + mem []byte // cache of pmem.Data() + + file1 File + file2 File + leaf File + closed bool + err error // sticky error +} + +// broken marks the tree broken with err as the reason. +// Any method on t or function taking a t as an argument +// is expected to call t.broken for I/O or data corruption errors. +// If the error comes from another method on t or function taking t as an argument, +// then that callee can be assumed to have called t.broken. +func (t *diskTree) broken(err error) error { + if t.err == nil { + t.err = err + } + return err +} + +// Create creates a new, empty on-disk [Tree] stored in the two named files. +// The files must not already exist, unless they are both os.DevNull, +// in which case the Tree is held only in memory. +func Create(file1, file2, disk string) (Tree, error) { + return open(file1, file2, disk, os.O_WRONLY|os.O_CREATE|os.O_EXCL, "create") +} + +// Open opens an on-disk [Tree] stored in the two named files. +// The files must have been created by a previous call to [Create]. +func Open(file1, file2, disk string) (Tree, error) { + return open(file1, file2, disk, os.O_RDWR, "open") +} + +func open(file1, file2, file3 string, mode int, op string) (Tree, error) { + f1, err := os.OpenFile(file1, mode, 0666) + if err != nil { + return nil, err + } + f2, err := os.OpenFile(file2, mode, 0666) + if err != nil { + f1.Close() + return nil, err + } + if op == "create" { + mode = os.O_RDWR | os.O_CREATE | os.O_EXCL + } + f3, err := os.OpenFile(file3, mode, 0666) + if err != nil { + f1.Close() + f2.Close() + return nil, err + } + return memOpen(f1, f2, f3, op) +} + +// New creates or opens an on-disk [Tree] in the given files. +// If both files are empty, New creates a new tree in those files. +// Otherwise, New opens a pre-existing tree stored in those files. +// Only one file contains the latest tree at a time, but the +// implementation alternates between files to implement atomic updates. +func New(file1, file2, file3 File) (Tree, error) { + var op string + var buf [1]byte + n1, err1 := file1.ReadAt(buf[:], 0) + n2, err2 := file2.ReadAt(buf[:], 0) + if n1 == 0 && n2 == 0 && err1 == io.EOF && err2 == io.EOF { + op = "create" + } else { + op = "open" + } + return memOpen(file1, file2, file3, op) +} + +// memOpen is the general implementation of open. +// op is "create", "open", or "new", indicating the operation +// being performed on the files; sync indicates whether to +// try to use the files' Sync method. +// (When using /dev/null for an in-memory tree, +// we avoid calling Sync, because it will fail.) +func memOpen(file1, file2, disk File, op string) (_ Tree, err error) { + pmemOp := pmem.Open + if op == "create" { + pmemOp = pmem.Create + } + mem, err := pmemOp("mpt tree\n", file1, file2, disk) + if err != nil { + return nil, err + } + t := &diskTree{ + pmem: mem, + file1: file1, + file2: file2, + leaf: disk, + } + defer func() { + if err != nil { + mem.Release() + mem.UnsafeUnmap() + } + }() + + runtime.AddCleanup(t, func(*struct{}) { mem.Release() }, nil) + + if op == "create" { + // Write initial tree. + mem, err := t.pmem.Expand(hdrSize) + if err != nil { + return nil, err + } + h := emptyTreeHash() + if err := t.mutate(mem[hdrHash:], h[:]); err != nil { + return nil, err + } + if err := t.pmem.Sync(); err != nil { + return nil, err + } + } + + t.mem = t.pmem.Data() + + return t, nil +} + +var errCorrupt = errors.New("corrupt tree data") + +// Sync syncs written data to disk. +func (t *diskTree) Sync() error { + t.mmu.RLock() + defer t.mmu.RUnlock() + + if t.err != nil { + return t.err + } + if !t.hdr().dirty() && !t.hdr().exact() { + if err := t.hdr().setExact(t, true); err != nil { + return err + } + } + if err := t.pmem.Sync(); err != nil { + return t.broken(err) + } + return nil +} + +// TODO figure out whether pmem should Close. + +// Close closes the tree and the files it uses. +func (t *diskTree) Close() error { + t.mmu.Lock() + defer t.mmu.Unlock() + + if t.closed { + return fmt.Errorf("tree already closed") + } + if err := t.pmem.Sync(); err != nil { + return t.broken(err) + } + if err := t.pmem.Release(); err != nil { + t.broken(err) + } + if err := t.pmem.UnsafeUnmap(); err != nil { + t.broken(err) + } + t.mem = nil + t.pmem = nil + if err := t.file1.Close(); err != nil { + t.broken(err) + } + if err := t.file2.Close(); err != nil { + t.broken(err) + } + if t.err != nil { + return t.err + } + t.closed = true + t.err = errors.New("tree is closed") // stop future method calls + return nil +} + +// TODO: should mutate be done by editing dst in place and then calling t.mutated(dst)? + +// mutate is like copy(dst, src) where dst is inside t.mem. +// It also records the mutation in the patch buffer, to be written +// to disk when the current patch block fills or Sync is called. +func (t *diskTree) mutate(dst, src []byte) error { + n := min(len(dst), len(src)) + if err := t.pmem.Mutate(dst[:n], src[:n]); err != nil { + return t.broken(err) + } + return nil +} + +// addrToMem returns the tree memory at address a and length n. +func (t *diskTree) addrToMem(a addr, n int) ([]byte, error) { + if a > addr(len(t.mem)) || len(t.mem)-int(a) < n { + return nil, t.broken(errCorrupt) + } + return t.mem[a : a+addr(n)], nil +} + +// memToAddr converts a byte slice p, which must be from t.mem, +// into an addr. +func (t *diskTree) memToAddr(p []byte) addr { + off, ok := t.pmem.Offset(p) + if !ok { + panic("mpt: memToAddr misuse") + } + return addr(off) +} + +// alloc allocates n more bytes of tree memory, returning it as a slice. +func (t *diskTree) alloc(n int) ([]byte, error) { + if cap(t.mem)-len(t.mem) < n { + mem, err := t.pmem.Expand(len(t.mem) + n) + if err != nil { + t.err = err + return nil, err + } + t.mem = mem[:len(t.mem)] + } + off := len(t.mem) + t.mem = t.mem[:off+n] + return t.mem[off : off+n], nil +} + +// An addr is an offset into the disk layout. +// It is stored on disk as a 48-bit big-endian value. +type addr uint64 + +// parseAddr returns the node address at the given byte offset. +func parseAddr(p []byte) addr { + return addr(binary.BigEndian.Uint16(p))<<32 | addr(binary.BigEndian.Uint32(p[2:])) +} + +// putAddr stores the node address at the given byte offset. +func putAddr(p []byte, a addr) { + binary.BigEndian.PutUint32(p[2:], uint32(a)) + binary.BigEndian.PutUint16(p, uint16(a>>32)) +} + +// A diskHdr is the memory copy of the tree header. +type diskHdr [hdrSize]byte + +func (h *diskHdr) version() int64 { return int64(binary.BigEndian.Uint64(h[hdrVersion:])) } +func (h *diskHdr) dirty() bool { return h[hdrDirty] != 0 } +func (h *diskHdr) exact() bool { return h[hdrExact] != 0 } +func (h *diskHdr) root() addr { return parseAddr(h[hdrRoot:]) } +func (h *diskHdr) hash() Hash { return Hash(h[hdrHash:]) } + +func (h *diskHdr) setVersion(t *diskTree, version int64) error { + var buf [8]byte + binary.BigEndian.PutUint64(buf[:], uint64(version)) + return t.mutate(h[hdrVersion:], buf[:]) +} + +func (h *diskHdr) setDirty(t *diskTree, d bool) error { + var buf [1]byte + if d { + buf[0] = 1 + } + return t.mutate(h[hdrDirty:], buf[:]) +} + +func (h *diskHdr) setExact(t *diskTree, d bool) error { + var buf [1]byte + if d { + buf[0] = 1 + } + return t.mutate(h[hdrExact:], buf[:]) +} + +func (h *diskHdr) setRoot(t *diskTree, n *diskNode) error { + a := t.addr(n) + var buf [6]byte + putAddr(buf[:], a) + return t.mutate(h[hdrRoot:], buf[:]) +} + +func (h *diskHdr) setHash(t *diskTree, hash Hash) error { + return t.mutate(h[hdrHash:], hash[:]) +} + +// hdr returns a pointer to the in-memory tree header. +func (t *diskTree) hdr() *diskHdr { + mem, err := t.addrToMem(0, hdrSize) + if err != nil { + panic(err) // mem should always be big enough for the header + } + return (*diskHdr)(mem) +} + +// A diskNode is the memory copy of a node. +// The *diskNodes passed around in this implementation +// are pointers into the in-memory copy t.mem. +type diskNode [nodeSize]byte + +// node returns the diskNode at the given address. +func (t *diskTree) node(a addr) (*diskNode, error) { + if a == 0 { + return nil, nil + } + mem, err := t.addrToMem(a, nodeSize) + if err != nil { + return nil, err + } + return (*diskNode)(mem), nil +} + +// addr returns the address of the given diskNode. +func (t *diskTree) addr(n *diskNode) addr { + if n == nil { + return 0 + } + return t.memToAddr(n[:]) +} + +// addrAt reads a node address from the address a. +// The caller must ensure that a is a valid address, +// or else addrAt panics. +func (t *diskTree) addrAt(a addr) addr { + mem, err := t.addrToMem(a, addrSize) + if err != nil { + panic(err) + } + return parseAddr(mem) +} + +// setAddrAt writes the node address b to the address a. +func (t *diskTree) setAddrAt(a, b addr) error { + mem, err := t.addrToMem(a, addrSize) + if err != nil { + return err + } + var buf [addrSize]byte + putAddr(buf[:], b) + return t.mutate(mem, buf[:]) +} + +// newNode allocates and returns a new node in the tree. +func (t *diskTree) newNode() (*diskNode, error) { + n, err := t.alloc(nodeSize) + if err != nil { + return nil, err + } + return (*diskNode)(n), nil +} + +func (n *diskNode) leafAddr(t *diskTree) int64 { + return int64(leafSize * ((t.addr(n) - hdrSize) / nodeSize)) +} + +func (n *diskNode) keyVal(t *diskTree) (Key, Val, error) { + var kv [leafSize]byte + if err := t.pmem.ReadDisk(kv[:], n.leafAddr(t)); err != nil { + return Key{}, Val{}, err + } + return Key(kv[:]), Val(kv[len(Key{}):]), nil +} + +func (n *diskNode) dirty() bool { return n[nodeDirty] != 0 } +func (n *diskNode) left() addr { return parseAddr(n[nodeLeft:]) } +func (n *diskNode) right() addr { return parseAddr(n[nodeRight:]) } +func (n *diskNode) ihash() Hash { return Hash(n[nodeIHash:]) } + +// bit returns the bit number recorded in the node. +// The single leaf node that is not also an inner node, +// identified by having no children, has bit number -1. +func (n *diskNode) bit() int { + if n.left() == 0 && n.right() == 0 { + return -1 + } + return int(n[nodeUbit]) +} + +// init initializes the node n with the given key, val, bit, left, and right; +// it also sets dirty=true and clears ihash. +func (n *diskNode) init(t *diskTree, key Key, val Val, bit int, left, right *diskNode) error { + var buf [nodeSize]byte + buf[nodeUbit] = byte(bit) + buf[nodeDirty] = 1 + putAddr(buf[nodeLeft:], t.addr(left)) + putAddr(buf[nodeRight:], t.addr(right)) + if err := t.mutate(n[:], buf[:]); err != nil { + return err + } + + var kv [leafSize]byte + copy(kv[:], key[:]) + copy(kv[len(Key{}):], val[:]) + if err := t.pmem.WriteDisk(kv[:], n.leafAddr(t)); err != nil { + return t.broken(err) + } + return nil +} + +func (n *diskNode) setVal(t *diskTree, val Val) error { + off := n.leafAddr(t) + int64(len(Key{})) + if err := t.pmem.WriteDisk(val[:], off); err != nil { + return t.broken(err) + } + return nil +} + +func (n *diskNode) setIHash(t *diskTree, h Hash) error { return t.mutate(n[nodeIHash:], h[:]) } +func (n *diskNode) setDirty(t *diskTree, d bool) error { + var p [1]byte + if d { + p[0] = 1 + } + return t.mutate(n[nodeDirty:], p[:]) +} diff --git a/mpt/disk_test.go b/mpt/disk_test.go new file mode 100644 index 0000000..c531ec1 --- /dev/null +++ b/mpt/disk_test.go @@ -0,0 +1,324 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package mpt + +import ( + "bytes" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "fmt" + "io" + "maps" + "math/rand/v2" + "runtime/debug" + "slices" + "testing" +) + +// A memFile is an in-memory file with ReadAt, WriteAt, Close, and Sync methods. +type memFile struct { + readOnly bool + data []byte +} + +func (f *memFile) ReadAt(data []byte, off int64) (int, error) { + if off < 0 || off >= int64(len(f.data)) { + return 0, io.EOF + } + n := copy(data, f.data[off:]) + if n < len(data) { + return n, io.ErrUnexpectedEOF + } + return n, nil +} + +func (f *memFile) WriteAt(data []byte, off int64) (int, error) { + if f.readOnly { + panic("write to read-only file") + } + if off > int64(len(f.data)) { + // Fill hole in file. + f.data = append(f.data, make([]byte, int(off)-len(f.data))...) + } + n := copy(f.data[off:], data) + f.data = append(f.data, data[n:]...) + return len(data), nil +} + +func (f *memFile) Close() error { + return nil +} + +func (f *memFile) Sync() error { + return nil +} + +func memHash(t *diskTree) string { + h := sha256.New() + h.Write(t.mem) + n := 1 + (len(t.mem)-hdrSize)/nodeSize + const pmemHdrSize = 16 + leaf := pmemHdrSize + n*64 + switch f := t.leaf.(type) { + default: + panic(fmt.Sprintf("unknown leaf type %T", t.leaf)) + case *memFile: + h.Write(f.data[:leaf]) + case *testFile: + if len(f.data) != leaf { + panic(fmt.Sprintf("unexpected leaf size in real tree: %d != %d (t.mem=%d)", len(f.data), leaf, len(t.mem))) + } + h.Write(f.data) + } + s := base64.StdEncoding.EncodeToString(h.Sum(nil)) + return fmt.Sprintf("%s/%#x", s[:7], len(t.mem)) +} + +// A tester is a two-file simulator that checks after each write that +// reopening the disk works properly, even if the write only happens +// partially or even gets corrupted (unlikely but we can handle it). +type tester struct { + t *testing.T + tree *diskTree // in-memory tree + file [3]testFile // files backing tree + valid map[string]bool // hashes of acceptable tree memory images + replay []int // replay log for recovery +} + +// A testFile is a single simulated file. +type testFile struct { + memFile + tester *tester +} + +func (f *testFile) name() string { + if f.tester == nil { + return "???" + } + for i := range 3 { + if f == &f.tester.file[i] { + return fmt.Sprint("file", i+1) + } + } + return "???" +} + +func (f *testFile) clone() *memFile { + return &memFile{readOnly: true, data: bytes.Clone(f.data)} +} + +// WriteAt writes to the test file. +func (f *testFile) WriteAt(data []byte, off int64) (int, error) { + f.tester.t.Logf("%s write %#x+%#x = %#x", f.name(), off, len(data), off+int64(len(data))) + return f.memFile.WriteAt(data, off) +} + +// Sync syncs the test file. +func (f *testFile) Sync() error { + if f.tester == nil { + panic("sync of read-only file") + } + + f.tester.t.Logf("%s sync at %#x", f.name(), len(f.data)) + return nil +} + +func (tt *tester) markOK() { + h := memHash(tt.tree) + tt.t.Logf("ok %v", h) + tt.valid[h] = true +} + +func (tt *tester) test(minVer int64, minExact bool) { + tt.try(&tt.file[0], minVer, minExact) + tt.try(&tt.file[1], minVer, minExact) +} + +// try tries reopening the files with various i/o problems. +func (tt *tester) try(f *testFile, minVer int64, minExact bool) { + if tt.tree == nil { + // Initial tree not created yet. + return + } + + // Test file with write actually succeeding. + tt.reopen(minVer, minExact, "as written") +} + +func (tt *tester) reopen(minVer int64, minExact bool, format string, args ...any) { + kind := fmt.Sprintf(format, args...) + f1 := tt.file[0].clone() + f2 := tt.file[1].clone() + f3 := tt.file[2].clone() + f3.readOnly = false + tree, err := New(f1, f2, f3) + if err != nil { + tt.t.Fatalf("reopen: %s: %v", kind, err) + } + defer tree.Close() + + version, exact := tree.Version() + if err != nil { + tt.t.Fatalf("reopen: %s: %v", kind, err) + } + if version < minVer || minExact != exact { + tt.t.Fatalf("reopen: %s: version = %d,%v, want ≥ %d,%v", kind, version, exact, minVer, minExact) + } + if !exact { + f1.readOnly = false + f2.readOnly = false + + // Find [-1, version] marking snapshot of recorded version. + i := 0 + if version > 0 { + for i < len(tt.replay) && (tt.replay[i] != -1 || int64(tt.replay[i+1]) != version) { + i += 2 + } + if i >= len(tt.replay) { + tt.t.Fatalf("reopen: %s: recover %d %v: cannot find version %d", kind, version, exact, version) + } + i += 2 + } + // Replay rest of log. + for ; i < len(tt.replay); i += 2 { + if tt.replay[i] == -1 { + if _, err := tree.Snap(int64(tt.replay[i+1])); err != nil { + tt.t.Fatalf("reopen: %s: Snap: %v", kind, err) + } + } else { + if err := tree.Set(Key(v(tt.replay[i])), v(tt.replay[i+1])); err != nil { + tt.t.Fatalf("reopen: %s: Set: %v", kind, err) + } + } + } + } + + h := memHash(tree.(*diskTree)) + if !tt.valid[h] { + tt.t.Fatalf("reopen (%d %d): %s: (%d %v): invalid hash %v want %v\n\n%s\nactual tree:\n%s\nrecovered tree:\n%s\nactual leaf:\n%s\nrecovered leaf (%v):\n%s", + len(tt.file[0].data), len(tt.file[1].data), kind, + version, exact, + h, slices.Sorted(maps.Keys(tt.valid)), + debug.Stack(), + hexDump(tt.tree.mem), + hexDump(tree.(*diskTree).mem), + hexDump(tt.tree.leaf.(*testFile).data), + tree.(*diskTree).leaf.(*memFile) == f3, + hexDump(tree.(*diskTree).leaf.(*memFile).data)) + } +} + +func hexDump(data []byte) string { + return hex.Dump(data[:min(len(data), 1024)]) +} + +// TODO maybe for testing enable a pmem mode that +// writes every mutation to a separate patch, +// and then reopen after every file write? + +func TestDiskRecovery(t *testing.T) { + for i := range 10 { + t.Run(fmt.Sprint(i), testDiskRecovery) + } +} + +func testDiskRecovery(t *testing.T) { + tt := &tester{t: t} + for i := range tt.file { + tt.file[i].tester = tt + } + + xtree, err := New(&tt.file[0], &tt.file[1], &tt.file[2]) + if err != nil { + t.Fatal(err) + } + tree := xtree.(*diskTree) + defer tree.Close() // relelase pmem on test failure + + tree.pmem.SetConstantFlushing(true) + tt.tree = tree + tt.valid = make(map[string]bool) + tt.markOK() + version := int64(0) + exact := false + syncVersion := version + syncExact := false + + for range 10 { + switch r := rand.N(10); r { + default: + i := rand.N(100) + j := rand.N(100) + t.Logf("set %d %d", i, j) + tt.replay = append(tt.replay, i, j) + check(t, tree.Set(Key(v(i)), v(j))) + exact = false + syncExact = false + tt.markOK() + tt.test(syncVersion, syncExact) + + case 0, 1: + version++ + exact = true + t.Logf("snap %d", version) + tt.replay = append(tt.replay, -1, int(version)) + _, err := tree.Snap(version) + check(t, err) + tt.markOK() + tt.test(syncVersion, syncExact) + fallthrough + + case 3: + t.Log("sync") + check(t, tree.Sync()) + _, exact = tree.Version() + syncVersion = version + syncExact = exact + clear(tt.valid) + tt.markOK() + tt.test(syncVersion, syncExact) + } + } + + check(t, tree.Close()) +} + +func TestDiskReopen(t *testing.T) { + // Test that very basic tree written to disk can be reopened, restored. + // Simulations are all well and good, but test real files a bit too. + dir := t.TempDir() + tree1, err := Create(dir+"/tree1", dir+"/tree2", dir+"/disk") + if err != nil { + t.Fatal(err) + } + check(t, err) + defer tree1.Close() + + for i := range 10 { + check(t, tree1.Set(Key(v(i)), v(i))) + } + + _, err = tree1.Snap(1) + check(t, err) + check(t, tree1.Sync()) + + tree2, err := Open(dir+"/tree1", dir+"/tree2", dir+"/disk") + check(t, err) + defer tree2.Close() + + if !bytes.Equal(tree1.(*diskTree).mem, tree2.(*diskTree).mem) { + t.Fatalf("tree memory differs\n\n%s\n\n%s", + hex.Dump(tree1.(*diskTree).mem[:1024]), + hex.Dump(tree2.(*diskTree).mem[:1024])) + } +} + +func check(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} diff --git a/mpt/dmem.go b/mpt/dmem.go new file mode 100644 index 0000000..84c1c31 --- /dev/null +++ b/mpt/dmem.go @@ -0,0 +1,446 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package mpt + +import "fmt" + +// hash returns the hash for the given tree node. +// pbit is the parent bit depth, controlling whether n is viewed as a leaf. +func (n *diskNode) hash(t *diskTree, pbit int) (Hash, error) { + if n.bit() <= pbit { + key, val, err := n.keyVal(t) + if err != nil { + return Hash{}, err + } + return hashLeaf(key, val), nil + } + return n.ihash(), nil +} + +var lazyHash = false + +// unhash marks n's hash invalid or recomputes it, +// depending on the [lazyHash] setting. +func (n *diskNode) unhash(t *diskTree, pbit int) error { + if !lazyHash { + _, err := n.rehash(t, pbit, true) + return err + } + if n.dirty() { + return nil + } + return n.setDirty(t, true) +} + +// rehash updates n.hash if needed and then returns it. +func (n *diskNode) rehash(t *diskTree, pbit int, force bool) (Hash, error) { + nbit := n.bit() + if nbit <= pbit { + return n.hash(t, pbit) + } + if n.dirty() || force { + left, err := t.node(n.left()) + if err != nil { + return Hash{}, err + } + lhash, err := left.rehash(t, nbit, false) + if err != nil { + return Hash{}, err + } + right, err := t.node(n.right()) + if err != nil { + return Hash{}, err + } + rhash, err := right.rehash(t, nbit, false) + if err != nil { + return Hash{}, err + } + if err := n.setIHash(t, hashInner(nbit, lhash, rhash)); err != nil { + return Hash{}, err + } + if err := n.setDirty(t, false); err != nil { + return Hash{}, err + } + } + return n.ihash(), nil +} + +// Snap returns a snapshot of t. +func (t *diskTree) Snap(version int64) (Snapshot, error) { + t.mmu.RLock() + defer t.mmu.RUnlock() + + if err := t.snap(version); err != nil { + return Snapshot{}, err + } + _ = t.check // t.check() + return Snapshot{t.hdr().version(), t.hdr().hash()}, nil +} + +func (t *diskTree) snap(version int64) error { + if t.err != nil { + return t.err + } + if t.hdr().dirty() { + // Note: Not using a mutation group because we might be + // updating arbitrarily many hashes during rehash. + // Without group, ordering matters: write hash before dirty + // and both before version. + root, err := t.node(t.hdr().root()) + if err != nil { + return err + } + hash, err := root.rehash(t, -1, false) + if err != nil { + return err + } + if err := t.hdr().setHash(t, hash); err != nil { + return err + } + if err := t.hdr().setDirty(t, false); err != nil { + return err + } + } + if version >= 0 { + if err := t.hdr().setVersion(t, version); err != nil { + return err + } + } + return nil +} + +// Version returns version information about the tree. +func (t *diskTree) Version() (version int64, exact bool) { + hdr := t.hdr() + return hdr.version(), hdr.exact() +} + +// Set sets the value associated with key to val. +func (t *diskTree) Set(key Key, val Val) error { + t.mmu.RLock() + defer t.mmu.RUnlock() + + if t.err != nil { + return t.err + } + + if t.hdr().exact() { + // Clear exact and flush to disk (in the memory files) + // before we make any writes to the disk leaf file, + // so that we know the disk leaf file may be ahead of the memory file. + if err := t.hdr().setExact(t, false); err != nil { + return err + } + if err := t.pmem.Sync(); err != nil { + return err + } + } + + // Keep all writes for this Set in the same group. + // We write one node and the dirty field for log N nodes, + // so it fits easily in the mutation group limit. + t.pmem.BeginGroup() + defer t.pmem.EndGroup() + + if !t.hdr().dirty() { + if err := t.hdr().setDirty(t, true); err != nil { + return err + } + } + if t.hdr().root() == 0 { + n, err := t.newNode() + if err != nil { + return err + } + n.init(t, key, val, 0, nil, nil) + if err := t.hdr().setRoot(t, n); err != nil { + return err + } + } else { + b, err := t.setChild(-1, hdrRoot, key, val) + if err != nil { + return err + } + if b >= 0 { + panic("bad add") + } + root, err := t.node(t.hdr().root()) + if err != nil { + return err + } + if err := root.unhash(t, -1); err != nil { + return err + } + } + _ = t.check // t.check() + return nil +} + +func (n *diskNode) set(t *diskTree, pbit int, key Key, val Val) (int, error) { + nbit := n.bit() + if nbit <= pbit { + // view n as leaf + nkey, _, err := n.keyVal(t) + if err != nil { + return 0, err + } + b := nkey.overlap(key) + if b == keyBits { + if err := n.setVal(t, val); err != nil { + return 0, err + } + return -1, nil + } + // Caller must create a node splitting at bit b. + return b, nil + } + + ptr := t.addr(n) + nodeLeft + if nbit >= 0 && key.bit(nbit) != 0 { + ptr = t.addr(n) + nodeRight + } + b, err := t.setChild(nbit, ptr, key, val) + if err != nil { + return 0, err + } + if b < 0 { + if err := n.unhash(t, pbit); err != nil { + return 0, err + } + } + return b, nil +} + +func (t *diskTree) setChild(nbit int, childp addr, key Key, val Val) (int, error) { + child, err := t.node(t.addrAt(childp)) + if err != nil { + return 0, err + } + b, err := child.set(t, nbit, key, val) + if err != nil { + return 0, err + } + if nbit < b { + n, err := t.newNode() + if err != nil { + return 0, err + } + var left, right *diskNode + if key.bit(b) == 0 { + left, right = n, child + } else { + left, right = child, n + } + n.init(t, key, val, b, left, right) + if err := t.setAddrAt(childp, t.addr(n)); err != nil { + return 0, err + } + b = -1 + } + return b, nil +} + +// Predict returns the hash of the tree that would result from +// applying the given changes (sorted by key) to the tree, +// without modifying the tree. +func (t *diskTree) Predict(changes []KeyVal) (Hash, error) { + t.mmu.RLock() + defer t.mmu.RUnlock() + + if t.err != nil { + return Hash{}, t.err + } + if t.hdr().dirty() { + return Hash{}, ErrModifiedTree + } + + s, list, err := t.predict([]node{}, t.hdr().root(), -1, changes) + if err != nil { + return Hash{}, err + } + for _, kv := range list { + s = reduce(append(s, node{prefix(kv.Key, 256), hashLeaf(kv.Key, kv.Val)})) + } + return hashStack(s), nil +} + +// predict calculates the edited tree hash for the subtree at address a. +func (t *diskTree) predict(s []node, a addr, pbit int, list []KeyVal) ([]node, []KeyVal, error) { + if a == 0 { + return s, list, nil + } + + n, err := t.node(a) + if err != nil { + return nil, nil, err + } + key, val, err := n.keyVal(t) + if err != nil { + return nil, nil, err + } + nbit := n.bit() + bits := nbit + if nbit <= pbit { + bits = 256 + } + pkey := prefix(key, bits) + + // Stack modifications before node. + for len(list) > 0 && prefix(list[0].Key, bits).compare(pkey) < 0 { + k, v := list[0].Key, list[0].Val + list = list[1:] + s = reduce(append(s, node{prefix(k, 256), hashLeaf(k, v)})) + } + + // Stack leaf node, possibly replaced. + if bits == 256 { + if len(list) > 0 && list[0].Key == key { + val = list[0].Val + list = list[1:] + } + s = reduce(append(s, node{pkey, hashLeaf(key, val)})) + return s, list, nil + } + + // Stack entire subtree, if no modifications inside it. + if len(list) == 0 || pkey.compare(prefix(list[0].Key, bits)) < 0 { + h, err := n.hash(t, pbit) + if err != nil { + return nil, nil, err + } + s = reduce(append(s, node{pkey, h})) + return s, list, nil + } + + // Otherwise, apply modifications within subtree. + s, list, err = t.predict(s, n.left(), nbit, list) + if err != nil { + return nil, nil, err + } + s, list, err = t.predict(s, n.right(), nbit, list) + if err != nil { + return nil, nil, err + } + return s, list, nil +} + +// Prove returns a proof of the presence or absence of key in t. +func (t *diskTree) Prove(key Key) (Proof, error) { + t.mmu.RLock() + defer t.mmu.RUnlock() + + if t.err != nil { + return nil, t.err + } + if t.hdr().dirty() { + return nil, ErrModifiedTree + } + root, err := t.node(t.hdr().root()) + if err != nil { + return nil, err + } + if root == nil { + return Proof(proofEmpty), nil + } + return root.prove(t, -1, key) +} + +func (n *diskNode) prove(t *diskTree, pbit int, key Key) (Proof, error) { + nbit := n.bit() + if nbit <= pbit { + // view n as leaf + nkey, nval, err := n.keyVal(t) + if err != nil { + return nil, err + } + var p Proof + if nkey == key { + p = Proof(proofConfirm) + } else { + p = append(Proof(proofDeny), nkey[:]...) + } + return append(p, nval[:]...), nil + } + + childAddr, sibAddr := n.left(), n.right() + if key.bit(nbit) == 1 { + childAddr, sibAddr = sibAddr, childAddr + } + child, err := t.node(childAddr) + if err != nil { + return nil, err + } + sib, err := t.node(sibAddr) + if err != nil { + return nil, err + } + sibHash, err := sib.hash(t, nbit) + if err != nil { + return nil, err + } + + p, err := child.prove(t, nbit, key) + if err != nil { + return nil, err + } + return append(append(p, byte(nbit)), sibHash[:]...), nil +} + +func (t *diskTree) check() { + println("check") + root, err := t.node(t.hdr().root()) + if err != nil { + panic(err) + } + if root == nil { + return + } + var sawNil bool + h := root.check(t, 1, -1, &sawNil) + if h != t.hdr().hash() && !t.hdr().dirty() { + fmt.Printf("have %v want %v\n", t.hdr().hash(), h) + panic("bad hash") + } + if !sawNil { + panic("lost nil") + } + println("check OK") +} + +func (n *diskNode) check(t *diskTree, depth, pbit int, sawNil *bool) Hash { + if n.bit() == -1 { + if *sawNil { + panic("multiple nils") + } + *sawNil = true + } + if n.bit() <= pbit { + // view as leaf + nkey, nval, err := n.keyVal(t) + if err != nil { + panic(err) + } + fmt.Printf("%*sleaf(%d) %#x %v %v %#x %#x %v dirty=%v\n", depth*2, "", n.bit(), t.addr(n), nkey, nval, n.left(), n.right(), hashLeaf(nkey, nval), n.dirty()) + return hashLeaf(nkey, nval) + } + fmt.Printf("%*s%d %#x %#x %#x %v dirty=%v\n", depth*2, "", n.bit(), t.addr(n), n.left(), n.right(), n.ihash(), n.dirty()) + + left, err := t.node(n.left()) + if err != nil { + panic(err) + } + right, err := t.node(n.right()) + if err != nil { + panic(err) + } + h := hashInner(n.bit(), + left.check(t, depth+1, n.bit(), sawNil), + right.check(t, depth+1, n.bit(), sawNil)) + if h != n.ihash() && !n.dirty() { + fmt.Printf("%*shave %v want %v\n", depth*2, "", n.ihash(), h) + panic("bad hash") + } + return h +} diff --git a/mpt/internal/pmem/pmem.go b/mpt/internal/pmem/pmem.go new file mode 100644 index 0000000..e63957e --- /dev/null +++ b/mpt/internal/pmem/pmem.go @@ -0,0 +1,1085 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package pmem implements persistent, transactional +// memory backed by on-disk files. +// This package only runs on 64-bit Unix architectures. +// (It could be made to run on Windows; +// it cannot be made to run on 32-bit systems, nor on Plan 9.) +// +// Each memory image is represented by a [Mem] +// and stored in a pair of on-disk files. +// See the [Mem] documentation for more details. +package pmem + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "errors" + "fmt" + "hash" + "io" + "log" + "strings" + "time" + + "filippo.io/torchwood/mpt/internal/slicemath" + "filippo.io/torchwood/mpt/internal/span" +) + +const verboseIO = false + +// File Layout +// +// magic [multiple of 8 bytes] +// frames... +// +// Frame Layout +// +// id [16] +// seq [8] +// len [8] +// data [len] +// sum [32] +// +// Patch frames are a sequence of mutations, each of the form: +// +// offset [v] +// len [v] +// data [len] +// +// where [v] denotes a uvarint-encoded int. + +const ( + hashSize = sha256.Size + + // field offsets + frameID = 0 + frameSeq = 16 + frameLen = 24 + + // frame header size + frameSize = 32 + + // total frame overhead, including final checksum + frameExtra = frameSize + hashSize +) + +// MaxGroupBytes is the maximum number of modified bytes +// that a mutation group can store: 1 MB. +// Writing the same bytes over and over counts toward the +// group limit each time. +const MaxGroupBytes = 1 << 20 + +// Logically, these are all constants, but we use variables +// so that they can be made smaller for testing. +var ( + // maxMem is the maximum memory size. + // We reserve the entire address space when opening the memory, + // so that it can be extended in place rather than needing to + // move the memory. + // 16 TB (44 bits) should be far more than enough for all practical uses, + // while leaving 19 bits of headroom, so that many different + // on-disk files can be used simultaneously. + maxMem = 16 << 40 + + // maxVarint is the maximum number of bytes needed for a varint + // encoding numbers up to maxMem. maxMem is 1<<44, so we + // need 45 bits, and at 7 bits per byte, that's 7 bytes. + maxVarint = 7 + + // maxPatch is the maximum size of a patch block, + // which must be able to hold a full mutation group. + // In the worst case, the framing of every data byte in + // a mutation might be framed by a maxVarint-byte offset + // and a 1-byte count. That's 1MB*8 = 8 MB. + // There needs to be headroom for the memory-length patch, + // so the minimum would be 8MB + 8 bytes, but we bump the + // patch block size to 16 MB instead. + maxPatch = 16 << 20 +) + +var errCorrupt = errors.New("corrupt input file") + +// File is the interface needed for on-disk storage. +type File interface { + io.ReaderAt + io.WriterAt + Sync() error +} + +// DevNull returns a file like the Unix /dev/null: it can be written but is always empty. +// Passing two DevNull files to New creates a Mem with no on-disk backing. +func DevNull() File { + return new(devNull) +} + +type devNull struct{} + +func (*devNull) ReadAt(b []byte, off int64) (int, error) { return 0, io.EOF } +func (*devNull) WriteAt(b []byte, off int64) (int, error) { return len(b), nil } + +func (*devNull) Sync() error { return nil } + +// A Mem represents a persistent memory backed by a pair of on-disk files. +// The files must implement the [File] interface ([os.File] is the usual implementation). +// The on-disk footprint is approximately six times the size of the memory image. +// +// At any moment, the current memory image is stored in one of the files, +// while the other holds an in-progress image. The two files swap meanings +// once the in-progress image has been fully written. Alternating between the +// two files provides consistency as well as atomicity of updates, +// including transactional grouping of writes. +// +// [Create] creates a new memory, and [Open] opens an existing one. +// +// The [Mem.Data] method provides access to a view of the memory, +// but modifications to it must be made using [Mem.Mutate], +// so that the changes can be logged to disk as well. +// The memory starts out with zero length and can be expanded +// to larger sizes using [Mem.Expand]. +// (Shrinking the memory is not implemented.) +// +// Mutations can be grouped into atomic transactions using +// [Mem.BeginGroup] and [Mem.EndGroup]. +// +// Calling [Mem.Sync] ensures that all modifications have been flushed +// to the underlying files, guaranteeing that a future [Open] will observe them. +// +// Calling [Mem.Close] closes the memory and leaves the mapping unreadable. +// Future accesses to the slice data returned by [Mem.Data] must be avoided. +// Those accesses will fault, meaning they crash the program unless +// [runtime/debug.SetPanicOnFault] has been used. +type Mem struct { + magic string + id [16]byte + tmp [frameExtra]byte + ptmp [2 * binary.MaxVarintLen64]byte + span *span.Span + mem []byte + patched int // length of “patched” section of memory + current *writer + next *writer + disk File // disk-only (not in memory) storage + diskOff int64 // offset where user writes begin + patch []byte + group int // group start in patch, or -1 if not in group + groupData int // total group data + err error + closed bool + compact compact + + constantFlushing bool + + syncHook func() + mutateHook func() +} + +// A reader is the state for reading an input file. +type reader struct { + file File + id [16]byte // id found in file + seq uint64 // file sequence number + off int64 // read offset in file + hash hash.Hash // hash of frame + memLen int // memory length + tmp [max(frameSize, 2*hashSize)]byte +} + +// A writer is the state for writing to an output file. +type writer struct { + file File + seq uint64 // file sequence number + off int64 // write offset in file + wrote bool // any writes since last sync? + + hash hash.Hash + tmp [max(hashSize, frameSize)]byte +} + +// A compact is the state for compaction. +type compact struct { + hash hash.Hash + off int + end int +} + +// broken marks the memory broken with err as the reason +// and returns err back to the caller. +// Only the first error is recorded. +// Any method on m that calls a non-m-method is expected to +// call m.broken if the method fails, so that I/O or data corruption +// errors permanently break all future uses of the memory. +func (m *Mem) broken(err error) error { + if m.err == nil { + m.err = err + } + return err +} + +// Create initializes a new memory stored in mem1, mem2, +// with disk-only storage in disk if non-nil. +// +// The magic string is recorded at the start of the file to +// distinguish different uses of persistent memory files. +// A future call to [Open] must pass the same magic string. +// The magic string must not contain any NUL (\x00) bytes. +// +// Create does not check that the files are empty, in order +// to support using pre-allocated disk files or raw disk partitions. +func Create(magic string, mem1, mem2, disk File) (*Mem, error) { + return create(magic, mem1, mem2, disk) +} + +// newMem allocates a new Mem for use by Create and Open. +func newMem(magic string) (*Mem, error) { + if strings.Contains(magic, "\x00") { + return nil, fmt.Errorf("magic %q must not contain NUL (\x00) byte", magic) + } + magic += "\x00\x00\x00\x00\x00\x00\x00"[:7&-len(magic)] // pad to 8 bytes + + sp, err := span.Reserve(maxMem) + if err != nil { + return nil, err + } + m := &Mem{ + magic: magic, + span: sp, + patch: make([]byte, 0, maxPatch), + group: -1, + compact: compact{ + hash: sha256.New(), + }, + } + return m, nil +} + +func newWriter(file File, seq uint64) *writer { + return &writer{ + file: file, + hash: sha256.New(), + seq: seq, + } +} + +// create implements Create but avoids exposing named results in the docs. +func create(magic string, file1, file2, disk File) (_ *Mem, err error) { + m, err := newMem(magic) + if err != nil { + return nil, err + } + defer func() { + if err != nil { + m.span.Release() + m.span.UnsafeUnmap() + } + }() + + rand.Read(m.id[:]) + m.current = newWriter(file1, 1) + m.next = newWriter(file2, 0) + + if err := m.writeEmptyTree(m.current); err != nil { + return nil, err + } + if err := m.writeEmptyTree(m.next); err != nil { + return nil, err + } + if disk != nil { + w := newWriter(disk, 0) + if err := m.writeEmptyTree(w); err != nil { + return nil, err + } + m.disk = disk + m.diskOff = w.off + } + return m, nil +} + +func (m *Mem) writeEmptyTree(w *writer) error { + if err := w.write([]byte(m.magic)); err != nil { + return m.broken(err) + } + if err := m.writeFrame(w, nil); err != nil { + return err + } + return m.sync(w) +} + +// SetConstantFlushing sets whether the memory should write every +// mutation to the files as quickly as possible. The setting defaults to false, +// in which case mutations are written to the files only when enough mutations +// have accumulated to fill a patch block or when Sync is called. +// Enabling constant flushing instead writes a separate patch block for +// every individual mutation outside a group, and for each group. +// This is an inefficient use of both file space and I/O bandwidth, +// but it can be useful for testing purposes to explore the full set of +// possible intermediate memory images that might be observed +// after a crash. +func (m *Mem) SetConstantFlushing(on bool) { + m.constantFlushing = on + if on { + m.flushPatch(false) + } +} + +func (w *writer) write(b []byte) error { + if err := w.writeAt(b, w.off); err != nil { + return err + } + w.off += int64(len(b)) + return nil +} + +func (w *writer) writeAt(b []byte, off int64) error { + start := time.Now() + _, err := w.file.WriteAt(b, off) + if verboseIO { + log.Printf("write %d %.6fs\n", len(b), time.Since(start).Seconds()) + } + if err != nil { + return err + } + w.wrote = true + return nil +} + +// Open opens the memory stored in the file pair mem1, mem2, +// with disk-only storage in disk if non-nil. +// The magic string must match the one used when the +// file pair was created with [Create]. +func Open(magic string, mem1, mem2, disk File) (*Mem, error) { + return open(magic, mem1, mem2, disk) +} + +// open implements Open but avoids exposing named results in the docs. +func open(magic string, file1, file2, disk File) (_ *Mem, err error) { + m, err := newMem(magic) + if err != nil { + return nil, err + } + defer func() { + if err != nil { + m.span.Release() + m.span.UnsafeUnmap() + } + }() + + // Peek at initial frame in both files. + r1, err := m.readStart(file1) + if err != nil { + return nil, err + } + r2, err := m.readStart(file2) + if err != nil { + return nil, err + } + if r1.id != r2.id { + return nil, fmt.Errorf("inconsistent pmem files: mismatched IDs") + } + if r1.seq == r2.seq { + return nil, fmt.Errorf("inconsistent pmem files: identical sequence numbers (%#x == %#x)", r1.seq, r2.seq) + } + if r1.seq < r2.seq { + r1, r2 = r2, r1 + } + if disk != nil { + rd, err := m.readStart(disk) + if err != nil { + return nil, err + } + if rd.id != r1.id { + return nil, fmt.Errorf("inconsistent pmem files: disk ID does not match mem ID") + } + m.disk = disk + m.diskOff = rd.off + hashSize + } + if err := m.readFile(r1); err != nil { + return nil, err + } + m.current = newWriter(r1.file, r1.seq) + m.current.off = r1.off + m.next = newWriter(r2.file, 0) + return m, nil +} + +// readStart reads the start of the file and returns a reader +// that can read the remainder of the file as well as +// the initial metadata observed at the start. +// The reader is positioned immediately after the magic string. +func (m *Mem) readStart(file File) (*reader, error) { + r := &reader{ + file: file, + hash: sha256.New(), + } + magic := make([]byte, len(m.magic)) + if err := m.read(r, magic); err != nil { + return nil, err + } + if string(magic) != m.magic { + return nil, m.broken(fmt.Errorf("bad magic: %q != %q", + strings.TrimRight(string(magic), "\x00"), + strings.TrimRight(m.magic, "\x00"))) + } + + var err error + r.id, r.seq, r.memLen, err = r.readFrameHeader() + if err != nil { + return nil, m.broken(err) + } + return r, nil +} + +// read reads exactly len(data) bytes into data from r.file at r.off. +func (m *Mem) read(r *reader, data []byte) error { + _, err := r.file.ReadAt(data, r.off) + if err != nil { + return m.broken(err) + } + r.off += int64(len(data)) + return nil +} + +// readFrameHeader reads and parses the next frame header from r. +// It resets r.hash and writes the header to it. +func (r *reader) readFrameHeader() (id [16]byte, seq uint64, n int, err error) { + f := r.tmp[:frameSize] + if _, err = r.file.ReadAt(f, r.off); err != nil { + return + } + r.off += int64(len(f)) + r.hash.Reset() + r.hash.Write(f) + copy(id[:], f[frameID:]) + seq = binary.BigEndian.Uint64(f[frameSeq:]) + n = int(binary.BigEndian.Uint64(f[frameLen:])) + if n < 0 { + err = errCorrupt + } + return +} + +// readFrame reads a single framed block from r, +// storing the data into data. +// If data is not large enough to hold the framed data, +// readFrame returns errCorrupt. +func (r *reader) readFrame(data []byte) (int, error) { + id, seq, n, err := r.readFrameHeader() + if err != nil { + return 0, err + } + if id != r.id || seq != r.seq || n > len(data) { + return 0, errCorrupt + } + if _, err := r.file.ReadAt(data[:n], r.off); err != nil { + return 0, err + } + r.off += int64(n) + r.hash.Write(data[:n]) + fsum := r.tmp[:hashSize] + if _, err := r.file.ReadAt(fsum, r.off); err != nil { + return 0, err + } + r.off += int64(len(fsum)) + hsum := r.hash.Sum(r.tmp[hashSize:hashSize]) + if [hashSize]byte(fsum) != [hashSize]byte(hsum) { + return 0, errCorrupt + } + return n, nil +} + +// readFile reads an entire memory image file from r. +// It assumes the magic string has been checked already. +// It reads an initial memory image of length memLen bytes +// followed by any number of patch blocks modifying or +// extending that image. +func (m *Mem) readFile(r *reader) error { + r.off = int64(len(m.magic)) + mem, err := m.span.Expand(r.memLen) + if err != nil { + return m.broken(err) + } + n, err := r.readFrame(mem[:r.memLen]) + if err != nil { + return m.broken(err) + } + if n != r.memLen { + // Unreachable unless file is changing underfoot. + // Caller just read the frame length and it was memLen. + return m.broken(errCorrupt) + } + m.mem = mem + m.patched = len(mem) + + patch := make([]byte, maxPatch) + for { + n, err := r.readFrame(patch) + if err == errCorrupt || err == io.EOF || err == io.ErrUnexpectedEOF { + break + } + if err != nil { + return m.broken(err) + } + if err := m.replay(patch[:n]); err != nil { + return err + } + } + return nil +} + +// replay applies the mutations listed in patch to m.mem. +// It expands m.mem as needed to apply the patch. +func (m *Mem) replay(patch []byte) error { + for len(patch) > 0 { + off, n := binary.Uvarint(patch) + if n <= 0 { + return m.broken(errCorrupt) + } + isDisk := off&1 != 0 + off >>= 1 + patch = patch[n:] + count, n := binary.Uvarint(patch) + if n <= 0 { + return m.broken(errCorrupt) + } + patch = patch[n:] + if count > uint64(len(patch)) || off+count < off || int(off+count) < 0 { + return m.broken(errCorrupt) + } + if isDisk { + // disk patch + if _, err := m.disk.WriteAt(patch[:count], m.diskOff+int64(off)); err != nil { + return m.broken(err) + } + } else { + // memory patch + if off+count > uint64(len(m.mem)) { + mem, err := m.span.Expand(int(off + count)) + if err != nil { + return m.broken(err) + } + m.mem = mem + } + copy(m.mem[off:off+count], patch[:count]) + m.patched = max(m.patched, int(off+count)) + } + patch = patch[count:] + } + return nil +} + +// Data returns the current memory. +// Changes to the memory must be made only using [Mem.Mutate], +// never using direct writes. Changes made by direct write will not +// be visible when the memory is reloaded by a future [Open]. +func (m *Mem) Data() []byte { + return m.mem +} + +// Expand extends the length of the current memory +// to be at least n bytes and returns the extended slice. +func (m *Mem) Expand(n int) ([]byte, error) { + if m.err != nil { + return nil, m.err + } + if n <= len(m.mem) { + return m.mem, nil + } + mem, err := m.span.Expand(n) + if err != nil { + // Do not use m.broken - nothing is broken yet. + // Caller might recover gracefully from being unable + // to expand the memory. + return nil, err + } + m.mem = mem + return m.mem, nil +} + +// Offset returns the starting offset of b within the memory. +// If b is not a subslice of m.Data(), Offset returns 0, false. +func (m *Mem) Offset(b []byte) (offset int, ok bool) { + off, ok := slicemath.Offset(m.mem, b) + return int(off), ok +} + +// BeginGroup starts an atomic mutation group. +// Expand and Mutate calls between Begin and [Mem.EndGroup] +// are guaranteed to be observed as an atomic unit +// upon reloading the memory: either they will all be +// present or none of them will be. +// Calls to BeginGroup must be followed eventually by a call to EndGroup +// and cannot be nested: it is an error to call BeginGroup twice +// without an intervening EndGroup. +// +// A group is limited to mutation of at most MaxGroupBytes bytes of mutated data. +func (m *Mem) BeginGroup() error { + if m.err != nil { + return m.err + } + if m.group >= 0 { + return fmt.Errorf("atomic mutation group already begun") + } + + // Patch buffer always has room to add an empty mutation + // at the end of the memory, to represent the most recent Expand. + // If the group grows too large, we will flush up to but not + // including the group, so add the empty mutation now. + if err := m.addMemLenPatch(); err != nil { + return err + } + + m.group = len(m.patch) + m.groupData = 0 + return nil +} + +// Mutate is like copy(dst, src), where dst must be inside m.Data() +// and src and dst must have the same length, +// but it also arranges to record the change on disk, +// so that it will be visible when the memory is reloaded. +func (m *Mem) Mutate(dst, src []byte) error { + if m.err != nil { + return m.err + } + if len(dst) != len(src) { + return fmt.Errorf("mismatched dst, src len in mutation") + } + if len(dst) == 0 { + return fmt.Errorf("empty mutation") + } + off, ok := m.Offset(dst) + if !ok { + return fmt.Errorf("invalid dst for mutation") + } + return m.mutate(uint64(off)<<1, src, func() error { + copy(dst, src) + m.patched = max(m.patched, off+len(src)) + return nil + }) +} + +// WriteDisk writes src to the disk-only file at offset off. +// It guarantees that on recovery after a crash, +// all disk writes that occurred before the latest recovered Mutate +// will be available for reading. +// (Disk writes that happened after that Mutate may or may not +// be available for reading as well.) +func (m *Mem) WriteDisk(src []byte, off int64) error { + if m.err != nil { + return m.err + } + if len(src) == 0 { + return nil + } + return m.mutate(uint64(off)<<1|1, src, func() error { + _, err := m.disk.WriteAt(src, m.diskOff+off) + if err != nil { + return m.broken(err) + } + return nil + }) +} + +// ReadDisk reads into dst from the disk-only file at offset off. +func (m *Mem) ReadDisk(dst []byte, off int64) error { + if m.err != nil { + return m.err + } + _, err := m.disk.ReadAt(dst, m.diskOff+off) + if err != nil { + if err == io.EOF { + err = io.ErrUnexpectedEOF + } + return m.broken(err) + } + return nil +} + +// mutate logs a write to the patch block, starting a new patch block if necessary. +// It calls commit to apply the actual write once it has checked a few +// error conditions. +// +// The offset off is the one recorded in the patch block. +// For mutations of the memory at offset o, off should be o<<1. +// For mutations of the disk-only file at offset o, off should be o<<1|1. +func (m *Mem) mutate(off uint64, src []byte, commit func() error) error { + // Note: it is tempting to return early if bytes.Equal(dst, src) is true, + // but there are two problems with that. One is that some callers + // may modify dst in place and then call m.Mutate(dst, dst). + // The other is that Mem.End depends on len(dst)==0 emitting + // a mutation, and of course all zero-length slices are equal. + + if m.group >= 0 { + if m.groupData+len(src) > MaxGroupBytes { + return fmt.Errorf("mutation group too large") + } + } + p := m.ptmp[:0] + p = binary.AppendUvarint(p, off) + p = binary.AppendUvarint(p, uint64(len(src))) + if len(m.patch)+len(p)+len(src)+maxVarint+1 > maxPatch { + if err := m.flushPatch(true); err != nil { + return err + } + } + m.patch = append(m.patch, p...) + m.patch = append(m.patch, src...) + if m.group >= 0 { + m.groupData += len(src) + } + if commit != nil { + if err := commit(); err != nil { + return err + } + } + if m.mutateHook != nil && m.group < 0 { + m.mutateHook() + } + if m.group < 0 && m.constantFlushing { + if err := m.flushPatch(true); err != nil { + return err + } + } + return nil +} + +// flushPatch flushes the current patch buffer to disk. +// If there is an active mutation group, only the buffer before +// that group is written. +func (m *Mem) flushPatch(needSpace bool) error { + var p []byte + if m.group >= 0 { + // Can only write up to m.group, but final mem len is already there. + if m.group == 0 { + if needSpace { + return m.broken(fmt.Errorf("pmem: internal error: group overflow")) + } + return nil + } + p = m.patch[:m.group] + } else { + if err := m.addMemLenPatch(); err != nil { + return err + } + p = m.patch + } + if len(p) == 0 { + return nil + } + + if err := m.writeFrame(m.current, p); err != nil { + return err + } + if m.next.seq != 0 { + if err := m.writeFrame(m.next, p); err != nil { + return err + } + } + m.patch = m.patch[:copy(m.patch, m.patch[len(p):])] // slide rest down + if m.group >= 0 { + m.group = 0 + } + return m.maybeCompact(2 * len(p)) +} + +// EndGroup finishes an atomic mutation group, +// which must have been started by [Mem.BeginGroup]. +func (m *Mem) EndGroup() error { + if m.err != nil { + return m.err + } + if m.group < 0 { + return fmt.Errorf("no atomic mutation group to end") + } + if m.patched != len(m.mem) { + // Before closing group, append an empty mutation if an Expand happened. + // Not using m.addMemLenPatch because we need to preserve the invariant + // that there will be room for _another_ when the eventual flush happens. + m.mutate(uint64(len(m.mem))<<1, nil, nil) + m.patched = len(m.mem) + } + m.group = -1 + + if m.next.seq > 0 && m.compact.off == m.compact.end && len(m.patch) > 0 { + m.flushPatch(false) + } + return nil +} + +// addMemLenPatch adds a final “memory length” patch to m.patch. +// Mutate ensures that there is always room for this final patch, +// so the error return should never happen. +func (m *Mem) addMemLenPatch() error { + if m.patched == len(m.mem) { + return nil + } + if len(m.patch)+maxVarint+1 > maxPatch { + return m.broken(fmt.Errorf("pmem internal patch overflow")) + } + m.patch = binary.AppendUvarint(m.patch, uint64(len(m.mem))<<1) + m.patch = binary.AppendUvarint(m.patch, 0) + m.patched = len(m.mem) + return nil +} + +// writeFrame writes a frame containing data to w. +func (m *Mem) writeFrame(w *writer, data []byte) error { + f := m.tmp[:frameSize] + copy(f[frameID:], m.id[:]) + binary.BigEndian.PutUint64(f[frameSeq:], w.seq) + binary.BigEndian.PutUint64(f[frameLen:], uint64(len(data))) + + w.hash.Reset() + w.hash.Write(f) + if err := w.write(f); err != nil { + return m.broken(err) + } + + w.hash.Write(data) + if err := w.write(data); err != nil { + return m.broken(err) + } + + sum := w.hash.Sum(w.tmp[:0]) + if err := w.write(sum); err != nil { + return m.broken(err) + } + + return nil +} + +// maybeCompact runs a bit of compaction if needed, +// limiting I/O to writing at most n data bytes plus some framing. +func (m *Mem) maybeCompact(n int) error { + if m.next.seq == 0 && m.current.off < 2*int64(len(m.mem)) { + // Current disk file is less than twice the tree memory. + // Not worth compacting yem. + return nil + } + + c := &m.compact + if m.next.seq == 0 { + if verboseIO { + log.Print("compact start") + } + // Start a new compaction. + // Record current tree size (but not content), + // so we know where patches should be written. + c.end = len(m.mem) + c.off = 0 + c.hash.Reset() + + m.next.off = int64(len(m.magic) + frameSize + c.end + 32) + m.next.seq = m.current.seq + 1 + + // Hash the correct frame header. + var frame [frameSize]byte + copy(frame[frameID:], m.id[:]) + binary.BigEndian.PutUint64(frame[frameSeq:], m.next.seq) + binary.BigEndian.PutUint64(frame[frameLen:], uint64(c.end)) + c.hash.Write(frame[:]) + + // But write seq=0 to disk for now, so that if we crash before finishing, + // the next Open will not try to use this file. + // We will write the correct sequence number once everything is on disk. + binary.BigEndian.PutUint64(frame[frameSeq:], 0) + if err := m.next.writeAt(frame[:], int64(len(m.magic))); err != nil { + return m.broken(err) + } + } + + // Write at most n bytes of data, both to c.hash and to m.next.file. + // + // Note: If compaction were running in parallel with writes, + // we could copy from c.mem racily into a buffer and then write + // the buffer to both the hash and the file. As long as they are + // consistent, any racy reads would not matter, since the writes + // we are racing against would be written in patch form, even if + // we didn't see them here. However, since we run compaction + // interleaved with other work, there should be no writes to c.mem, + // and we can read from it twice. + if c.off < c.end && n > 0 { + n := min(n, c.end-c.off) + c.hash.Write(m.mem[c.off : c.off+n]) + if err := m.next.writeAt(m.mem[c.off:c.off+n], int64(len(m.magic)+frameSize+c.off)); err != nil { + return m.broken(err) + } + c.off += n + } + + if c.off < c.end || m.group >= 0 || len(m.patch) > 0 { + // Not finished. Wait for next call. + // Note that if c.off == c.end but m.group >= 0, + // then there is an active mutation group, and the memory image + // we wrote may include writes from that group. + // Similarly, if len(m.patch) > 0, the group may have ended + // but the patches have not yet been flushed + // (EndGroup will flush them for us but hasn't yet). + // We cannot complete the image until the group is flushed. + return nil + } + + // Wrote entire tree image. Finish and switch. + sum := c.hash.Sum(nil) + if err := m.next.writeAt(sum[:], int64(len(m.magic)+frameSize+c.off)); err != nil { + return m.broken(err) + } + + // Sync m.disk to disk, because we are about to abandon + // all the disk patches in m.current. + if m.disk != nil { + if err := m.disk.Sync(); err != nil { + return m.broken(err) + } + } + + // Open will start using the tree when the bigger sequence number hits the disk, + // so we want to make sure that happens last. + // Sync entire tree to disk, then update sequence number, then sync again. + if err := m.sync(m.next); err != nil { + return err + } + if err := m.writeFrameSeq(m.next, int64(len(m.magic)), m.next.seq); err != nil { + return err + } + if err := m.sync(m.next); err != nil { + return err + } + + if verboseIO { + log.Print("compact switch") + } + // Switch current and next. + m.current, m.next = m.next, m.current + setCurrent(m.current.file, true, int(m.current.off)) + setCurrent(m.next.file, false, int(m.next.off)) + m.next.seq = 0 + return nil +} + +// writeFrameSeq updates a frame header at the given offset, +// replacing the sequence number with seq and leaving the +// rest of the frame header unmodified. +func (m *Mem) writeFrameSeq(w *writer, off int64, seq uint64) error { + binary.BigEndian.PutUint64(m.tmp[:], seq) + if err := w.writeAt(m.tmp[:8], off+frameSeq); err != nil { + return m.broken(err) + } + return nil +} + +func setCurrent(f File, b bool, off int) { + if f, ok := f.(interface{ setCurrent(bool, int) }); ok { + f.setCurrent(b, off) + } +} + +// Sync flushes and syncs all memory changes to the underlying files. +// +// As changes are made with Mutate, they are flushed to disk +// incrementally, so that the in-memory footprint of a Mem +// is only a limited amount more than its memory data. +// Sync makes sure that all mutations have been written +// to the files and then calls [File.Sync] to sync those writes. +func (m *Mem) Sync() error { + if m.err != nil { + return m.err + } + if err := m.flushPatch(false); err != nil { + return err + } + if m.syncHook != nil { + m.syncHook() + } + if err := m.sync(m.current); err != nil { + return err + } + if m.next.seq != 0 { + if err := m.sync(m.next); err != nil { + return err + } + } + return nil +} + +func (m *Mem) sync(w *writer) error { + if w.wrote { + start := time.Now() + err := w.file.Sync() + if verboseIO { + log.Printf("sync %.6fs\n", time.Since(start).Seconds()) + } + if err != nil { + return m.broken(err) + } + w.wrote = false + } + return nil +} + +// Release syncs the memory and makes the in-memory data unreadable. +// Future accesses to the slice data will fault, causing the program to crash, +// unless [runtime/debug.SetPanicOnFault] has changed the fault behavior. +// +// Release releases the Mem's physical memory back to the operating system, +// so that it can be used for other purposes. +// However, to make Release a safe operation, Release preserves the virtual +// address space reservation, which only costs a few kilobytes to maintain +// until the process exits. +// To release the virtual address space, see [Mem.UnsafeUnmap], +// but read and understand the warnings in its doc comment before using it. +func (m *Mem) Release() error { + // Collect errors as we go, but make sure to reach end. No early returns. + m.Sync() + if m.mem != nil { + if err := m.span.Release(); err != nil { + m.broken(err) + } + m.mem = nil + } + m.closed = true + if m.err != nil { + return m.err + } + m.err = errors.New("mem already closed") // for next time + return nil +} + +// UnsafeUnmap unmaps the virtual address space used by a closed memory. +// If m has not been closed using [Mem.Close], UnsafeUnmap does nothing +// but return an error. +// +// Normally, calling [Mem.Close] is sufficient to release the Mem's resources, +// and UnsafeUnmap need not be used. The only reason to use UnsafeUnmap +// is because the program opens and closes hundreds of thousands of Mems +// and must unmap old ones to avoid running out of virtual address space. +// Programs that use only tens of thousands of Mems, or just a few, +// need not use UnsafeUnmap. +// +// After Close, future accesses to the slice data previously returned by [Mem.Data] +// are guaranteed to fault. After UnsafeUnmap, accesses may still fault, +// but if the operating system has reused the virtual address space for other +// purposes the accesses may succeed and read unrelated memory. +// This is why the method is considered unsafe. +// (Writes to the slice data would write that unrelated memory as well, +// but in a correct program there should not be any such writes, +// since writes should only be done using [Mem.Mutate].) +func (m *Mem) UnsafeUnmap() error { + if !m.closed { + return fmt.Errorf("mem not closed; cannot unmap") + } + m.span.UnsafeUnmap() + return nil +} + +// hash returns a short hash of the current memory content, +// useful for debugging and testing. +func (m *Mem) hash() string { + h := sha256.Sum256(m.mem) + s := base64.StdEncoding.EncodeToString(h[:]) + return fmt.Sprintf("%s/%#x", s[:7], len(m.mem)) +} diff --git a/mpt/internal/pmem/pmem_test.go b/mpt/internal/pmem/pmem_test.go new file mode 100644 index 0000000..924559f --- /dev/null +++ b/mpt/internal/pmem/pmem_test.go @@ -0,0 +1,282 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// TODO test constant flushing mode +// TODO test recovery of disk writes + +package pmem + +import ( + "bytes" + "encoding/hex" + "fmt" + "io" + "math/rand/v2" + "runtime/debug" + "testing" +) + +func TestRecovery(t *testing.T) { + for i := range 10 { + t.Run(fmt.Sprint(i), testRecovery) + } +} + +func testRecovery(t *testing.T) { + tmp := make([]byte, 1000) + + oldPatch := maxPatch + oldMem := maxMem + defer func() { + maxPatch = oldPatch + maxMem = oldMem + }() + + maxPatch = 256 + maxMem = 1 << 30 + + tt := &tester{t: t} + for i := range tt.file { + tt.file[i].tester = tt + } + + mem, err := Create("magic", &tt.file[0], &tt.file[1], nil) + if err != nil { + t.Fatal(err) + } + tt.setMem(mem) + tt.markOK() + + const ( + MaxOff = 100 + MaxCount = 100 + ) + for range 1000 { + switch rand.N(10) { + case 0, 1, 2, 3, 4: + // Write many random memory sections, + // more than will fit in a single patch block. + for range 5 { + off := rand.N(MaxOff) + n := 1 + rand.N(MaxCount) + tt.t.Logf("mutate %#x+%#x", off, n) + _, err := mem.Expand(off + n) + tt.markOK() + check(tt.t, err) + check(tt.t, mem.Mutate(mem.Data()[off:off+n], randFill(tmp[:n]))) + } + + case 5, 6, 7, 8: + // Write a pair of grouped updates. + // Have to limit to single patch block but try to use + // almost the entire block so that a flush will be needed. + n := maxPatch - 4*(maxVarint+1) + n1 := 1 + rand.N(n-1) + n2 := n - n1 + off1 := rand.N(MaxOff) + off2 := rand.N(MaxOff) + tt.t.Logf("begingroup (len=%#x)", len(mem.mem)) + check(tt.t, mem.BeginGroup()) + _, err := mem.Expand(off1 + n1) + check(tt.t, err) + tt.t.Logf("mutate %#x+%#x", off1, n1) + check(tt.t, mem.Mutate(mem.Data()[off1:off1+n1], randFill(tmp[:n1]))) + _, err = mem.Expand(off2 + n2) + check(tt.t, err) + tt.t.Logf("mutate %#x+%#x", off2, n2) + check(tt.t, mem.Mutate(mem.Data()[off2:off2+n2], randFill(tmp[:n2]))) + tt.t.Logf("endgroup") + tt.markOK() + check(tt.t, mem.EndGroup()) + + case 9: + // Sync. + tt.t.Logf("sync") + check(tt.t, mem.Sync()) + } + } + + check(t, mem.Release()) + check(t, mem.UnsafeUnmap()) +} + +func randFill(b []byte) []byte { + for i := range b { + b[i] = byte(rand.N(256)) + } + return b +} + +type tester struct { + t *testing.T + mem *Mem + file [2]testFile + valid map[string]bool // hashes of acceptable memory images +} + +type testFile struct { + tester *tester + data []byte // data in file + sync int // offset of last sync; writes only append + current bool // whether file is current +} + +func (f *testFile) setCurrent(current bool, off int) { + f.current = current + f.data = f.data[:off] +} + +func (f *testFile) clone() *testFile { + return &testFile{data: bytes.Clone(f.data)} +} + +// ReadAt reads from the test file. +func (f *testFile) ReadAt(data []byte, off int64) (int, error) { + if off < 0 || off >= int64(len(f.data)) { + return 0, io.EOF + } + n := copy(data, f.data[off:]) + if n < len(data) { + return n, io.ErrUnexpectedEOF + } + return n, nil +} + +// WriteAt writes to the test file. +func (f *testFile) WriteAt(data []byte, off int64) (int, error) { + if f.tester == nil { + panic("write to read-only file") + } + + // Writes to the current file should only ever append; + // not overwriting is part of our reliability story. + // Writes to the next file can be scattered, because + // we are writing the tree interleaved with new patches. + if f.current && off != int64(len(f.data)) { + return 0, fmt.Errorf("non-appending write\n\n%s", debug.Stack()) + } + if off > int64(len(f.data)) { + // Fill hole in file. + f.data = append(f.data, make([]byte, int(off)-len(f.data))...) + } + f.tester.t.Logf("%s write %#x+%#x = %#x", f.name(), off, len(data), off+int64(len(data))) + n := copy(f.data[off:], data) + f.data = append(f.data, data[n:]...) + + // Try corrupting the writes and see what happens. + f.tester.try(f) + + return len(data), nil +} + +// Close closes the test file. +func (f *testFile) Close() error { + return nil +} + +func (f *testFile) name() string { + if f.tester == nil { + return "???" + } + if f == &f.tester.file[0] { + return "file0" + } + return "file1" +} + +// Sync syncs the test file. +// After Sync, bytes before the current offset cannot be lost or corrupted. +func (f *testFile) Sync() error { + if f.tester == nil { + return nil + } + + f.sync = len(f.data) + f.tester.t.Logf("%s sync at %#x", f.name(), f.sync) + f.tester.try(f) + return nil +} + +func (tt *tester) setMem(mem *Mem) { + tt.mem = mem + mem.syncHook = tt.syncHook + mem.mutateHook = tt.markOK + if tt.valid == nil { + tt.valid = make(map[string]bool) + } + h := tt.mem.hash() + tt.t.Logf("initial hash %v", h) + tt.valid[h] = true +} + +func (tt *tester) markOK() { + tt.t.Helper() + h := tt.mem.hash() + tt.t.Logf("ok %s", h) + tt.valid[h] = true +} + +func (tt *tester) syncHook() { + clear(tt.valid) // older snapshots no longer acceptable + tt.markOK() +} + +// try tries reopening the files with various i/o problems. +func (tt *tester) try(f *testFile) { + if tt.mem == nil { + // Initial tree not created yet. + return + } + + tt.reopen("as written") + + // Test file truncated to last sync. + whole := f.data + f.data = whole[:f.sync] + tt.reopen("truncated to last sync at %#x", f.sync) + + // Test file truncated past the sync point. + if n := len(whole) - f.sync; n >= 2 { + for range 5 { + pos := f.sync + 1 + rand.N(n-1) + f.data = whole[:pos] + tt.reopen("truncated to %#x", pos) + } + } + + // Test file with correct length but corrupt data past the sync point. + f.data = whole + if len(f.data) > f.sync { + for range 5 { + pos := f.sync + rand.N(len(f.data)-f.sync) + f.data[pos] ^= 1 + tt.reopen("corrupted at %#x", pos) + f.data[pos] ^= 1 + } + } + + // Test file with write actually succeeding. + tt.reopen("as written") +} + +func (tt *tester) reopen(format string, args ...any) { + kind := fmt.Sprintf(format, args...) + mem, err := Open("magic", tt.file[0].clone(), tt.file[1].clone(), nil) + if err != nil { + tt.t.Fatalf("reopen: %s: %v\n\n%s", kind, err, hex.Dump(tt.file[0].data)) + } + h := mem.hash() + if !tt.valid[h] { + tt.t.Fatalf("reopen (%d %d): %s: invalid hash %v want %v\n\n%s\n\n%s\n\n%s", len(tt.file[0].data), len(tt.file[1].data), kind, h, tt.valid, debug.Stack(), hex.Dump(tt.mem.mem), hex.Dump(mem.mem)) + } + check(tt.t, mem.Release()) + check(tt.t, mem.UnsafeUnmap()) +} + +func check(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} diff --git a/mpt/internal/slicemath/slicemath.go b/mpt/internal/slicemath/slicemath.go new file mode 100644 index 0000000..a1a8d94 --- /dev/null +++ b/mpt/internal/slicemath/slicemath.go @@ -0,0 +1,23 @@ +// Package slicemath implements safe “pointer arithmetic” on slices. +// It is a separate package so that packages using it do not need to +// import unsafe directly. +package slicemath + +import "unsafe" + +// contains reports whether big contains little; +// that is, it reports whether little is a subslice of big. +func contains(big, little []byte) bool { + return uintptr(unsafe.Pointer(&big[0])) <= uintptr(unsafe.Pointer(&little[0])) && + uintptr(unsafe.Pointer(&little[len(little)-1])) <= uintptr(unsafe.Pointer(&big[len(big)-1])) +} + +// Offset reports little's starting position within big. +// If big does not contain little, Offset returns ^uintptr(0), false. +// The caller must have checked sliceContains(big, little) already. +func Offset(big, little []byte) (offset uintptr, ok bool) { + if !contains(big, little) { + return ^uintptr(0), false + } + return uintptr(unsafe.Pointer(&little[0])) - uintptr(unsafe.Pointer(&big[0])), true +} diff --git a/mpt/internal/span/span_unix.go b/mpt/internal/span/span_unix.go new file mode 100644 index 0000000..29fbd6a --- /dev/null +++ b/mpt/internal/span/span_unix.go @@ -0,0 +1,104 @@ +//go:build !plan9 && !windows + +// Package span implements growable memory spans. +package span + +import ( + "fmt" + + "golang.org/x/sys/unix" +) + +type Span struct { + max int + alloc int + mem []byte + released []byte +} + +const pageSize = 4 << 20 + +func round(n int) int { + return (n + pageSize - 1) &^ (pageSize - 1) +} + +// Reserve returns a Span with zero memory footprint +// but with space reserved to expand to at most max bytes. +func Reserve(max int) (*Span, error) { + r := round(max) + mem, err := unix.Mmap(-1, 0, r, unix.PROT_NONE, unix.MAP_ANON|unix.MAP_PRIVATE|unix.MAP_NORESERVE) + if err != nil { + return nil, fmt.Errorf("span.Reserve %d: %w", r, err) + } + return &Span{max: max, mem: mem}, nil +} + +// Expand expands the accessible memory to at least n bytes +// and returns a slice of length n and capacity n. +// Calling Expand with a small n does not release memory +// from an earlier call with a larger n. +func (s *Span) Expand(n int) ([]byte, error) { + r := round(n) + if s.alloc < r { + err := unix.Mprotect(s.mem[s.alloc:r], unix.PROT_READ|unix.PROT_WRITE) + if err != nil { + return nil, fmt.Errorf("span.Expand %d..%d: %w", s.alloc, r, err) + } + s.alloc = r + } + return s.mem[:n:n], nil +} + +// Release releases the memory for the span. +// If previously returned memory is accessed after calling Release, +// the accesses will fault, which will crash the program +// or else panic, depending on the use of [runtime/debug.SetPanicOnFault]. +// [Span.UnsafeUnmap] releases the virtual memory for the span, +// but it is unsafe and rarely necessary to use. +func (s *Span) Release() error { + if s.mem == nil { + return fmt.Errorf("span.Release already called") + } + if s.alloc > 0 { + if err := unix.Mprotect(s.mem[:s.alloc], unix.PROT_NONE); err != nil { + return fmt.Errorf("span.Release: mprotect: %w", err) + } + err := unix.Madvise(s.mem[:s.alloc], unix.MADV_FREE) + if err == unix.EINVAL { + // MADV_FREE is missing before Linux 4.5 and in gVisor. + err = unix.Madvise(s.mem[:s.alloc], unix.MADV_DONTNEED) + } + if err != nil { + return fmt.Errorf("span.Release: madvise: %w", err) + } + } + s.released = s.mem + s.mem = nil + return nil +} + +// UnsafeUnmap releases the virtual memory for the span, +// making accesses to the previously returned memory behave unpredictably. +// Perhaps they will still fault, but if the operating system reuses the +// virtual address space, they might instead access unrelated memory. +// On 64-bit systems, the virtual address space available to processes +// is typically on the order of 2⁶³ bytes. +// Unless [Reserve] is being called for sizes totaling beyond that amount, +// programs can use [Span.Release] without UnsafeUnmap and avoid +// the unsafe behavior. +// +// If [Span.Release] has not been called, UnsafeUnmap does nothing +// but return an error. +func (s *Span) UnsafeUnmap() error { + if s.mem != nil { + return fmt.Errorf("Span.UnsafeUnmap without Span.Release") + } + if s.released == nil { + return fmt.Errorf("Span.UnsafeUnmap already called") + } + if err := unix.Munmap(s.released); err != nil { + return fmt.Errorf("span.Release: %w", err) + } + s.released = nil + return nil +} diff --git a/mpt/mem.go b/mpt/mem.go new file mode 100644 index 0000000..9fea186 --- /dev/null +++ b/mpt/mem.go @@ -0,0 +1,307 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package mpt + +import ( + "errors" + "fmt" +) + +// A memTree is an in-memory [Tree]. +type memTree struct { + version int64 // version number of tree + exact bool // version is exact + root *memNode // root node + hash Hash // overall tree hash + dirty bool // Set called without Snap + err error // sticky error condition +} + +// A memNode is a single node in the in-memory tree. +type memNode struct { + key Key + val Val + ihash Hash + dirty bool // needs rehashing + ubit byte + left *memNode + right *memNode +} + +func (n *memNode) bit() int { + if n.left == nil && n.right == nil { + return -1 + } + return int(n.ubit) +} + +// NewMemTree returns a new in-memory [Tree]. +func NewMemTree() Tree { + t := &memTree{ + hash: emptyTreeHash(), + exact: true, + } + return t +} + +// hash returns the hash for the given tree node. +// pbit is the parent bit depth, controlling whether n is viewed as a leaf. +func (n *memNode) hash(pbit int) Hash { + if n.bit() <= pbit { + return hashLeaf(n.key, n.val) + } + return n.ihash +} + +// unhash marks n's hash invalid. +func (n *memNode) unhash() { + n.dirty = true +} + +// rehash updates n.hash if needed and then returns it. +func (n *memNode) rehash(pbit int) Hash { + nbit := n.bit() + if nbit <= pbit { + return hashLeaf(n.key, n.val) + } + if n.dirty { + n.ihash = hashInner(nbit, n.left.rehash(nbit), n.right.rehash(nbit)) + n.dirty = false + } + return n.ihash +} + +// Sync is a no-op since the data is only in memory. +func (t *memTree) Sync() error { + return nil +} + +// Close is a no-op since the data is only in memory. +func (t *memTree) Close() error { + if t.err != nil { + return t.err + } + t.err = errors.New("tree is closed") + return nil +} + +func (t *memTree) UnsafeUnmap() error { return nil } + +// Stat returns the tree metadata. +func (t *memTree) Version() (version int64, exact bool) { + return t.version, t.exact +} + +// Snap returns a snapshot of t. +func (t *memTree) Snap(version int64) (Snapshot, error) { + if t.err != nil { + return Snapshot{}, t.err + } + if t.dirty { + // nothing, but keep the read for causing races with Set + } + t.dirty = false + if version >= 0 { + t.version = version + } + if t.root != nil { + t.hash = t.root.rehash(-1) + _ = t.check // t.check() + } + t.exact = true + return Snapshot{t.version, t.hash}, nil +} + +// Set sets the value associated with key to val. +func (t *memTree) Set(key Key, val Val) error { + if t.err != nil { + return t.err + } + t.dirty = true + t.exact = false + if t.root == nil { + t.root = &memNode{key: key, val: val} + } else { + if setChild(-1, &t.root, key, val) >= 0 { + panic("bad add") + } + } + _ = t.check // t.check() + return nil +} + +func (n *memNode) set(pbit int, key Key, val Val) int { + if n.bit() <= pbit { + // view n as leaf + b := n.key.overlap(key) + if b == keyBits { + n.val = val + return -1 + } + // Caller must create a node splitting at bit b. + return b + } + + nbit := n.bit() + ptr := &n.left + if nbit >= 0 && key.bit(nbit) != 0 { + ptr = &n.right + } + b := setChild(nbit, ptr, key, val) + if b < 0 { + n.unhash() + } + return b +} + +func setChild(nbit int, child **memNode, key Key, val Val) int { + b := (*child).set(nbit, key, val) + if nbit < b { + n := new(memNode) + var left, right *memNode + if key.bit(b) == 0 { + left, right = n, *child + } else { + left, right = *child, n + } + *n = memNode{ + key: key, + val: val, + ubit: uint8(b), + dirty: true, + left: left, + right: right, + } + *child = n + b = -1 + } + return b +} + +// Predict returns the hash of the tree that would result from +// applying the given changes (sorted by key) to the tree, +// without modifying the tree. +func (t *memTree) Predict(changes []KeyVal) (Hash, error) { + if t.err != nil { + return Hash{}, t.err + } + if t.dirty { + return Hash{}, ErrModifiedTree + } + + s, list := t.predict([]node{}, t.root, -1, changes) + for _, kv := range list { + s = reduce(append(s, node{prefix(kv.Key, 256), hashLeaf(kv.Key, kv.Val)})) + } + return hashStack(s), nil +} + +func (t *memTree) predict(s []node, n *memNode, pbit int, list []KeyVal) ([]node, []KeyVal) { + if n == nil { + return s, list + } + key, val := n.key, n.val + nbit := n.bit() + bits := nbit + if nbit <= pbit { + bits = 256 + } + pkey := prefix(key, bits) + + // Stack modifications before node. + for len(list) > 0 && prefix(list[0].Key, bits).compare(pkey) < 0 { + k, v := list[0].Key, list[0].Val + list = list[1:] + s = reduce(append(s, node{prefix(k, 256), hashLeaf(k, v)})) + } + + // Stack leaf node, possibly replaced. + if bits == 256 { + if len(list) > 0 && list[0].Key == key { + val = list[0].Val + list = list[1:] + } + s = reduce(append(s, node{pkey, hashLeaf(key, val)})) + return s, list + } + + // Stack entire subtree, if no modifications inside it. + if len(list) == 0 || pkey.compare(prefix(list[0].Key, bits)) < 0 { + h := n.hash(pbit) + s = reduce(append(s, node{pkey, h})) + return s, list + } + + // Otherwise, apply modifications within subtree. + s, list = t.predict(s, n.left, nbit, list) + s, list = t.predict(s, n.right, nbit, list) + return s, list +} + +// Prove returns a proof of the presence or absence of key in t. +func (t *memTree) Prove(key Key) (Proof, error) { + if t.err != nil { + return nil, t.err + } + if t.dirty { + return nil, ErrModifiedTree + } + if t.root == nil { + return Proof(proofEmpty), nil + } + return t.root.prove(-1, key), nil +} + +func (n *memNode) prove(pbit int, key Key) Proof { + nbit := n.bit() + if nbit <= pbit { + // view n as leaf + var p Proof + if n.key == key { + p = Proof(proofConfirm) + } else { + p = append(Proof(proofDeny), n.key[:]...) + } + return append(p, n.val[:]...) + } + + var sib Hash + var child *memNode + if key.bit(nbit) == 0 { + child = n.left + sib = n.right.hash(nbit) + } else { + child = n.right + sib = n.left.hash(nbit) + } + return append(append(child.prove(nbit, key), byte(nbit)), sib[:]...) +} + +// check checks all the tree invariants, walking the entire tree. +// It is too slow for real use but helpful to insert when debugging. +func (t *memTree) check() { + println("check") + h := t.root.check(1, -1) + if h != t.hash && (t.root == nil || !t.dirty) { + fmt.Printf("have %v want %v\n", t.hash, h) + panic("bad hash") + } +} + +func (n *memNode) check(depth, pbit int) Hash { + nbit := n.bit() + if nbit <= pbit { + // view as leaf + fmt.Printf("%*s%d leaf %v %v %p %p %p %v\n", depth*2, "", n.bit(), n.key, n.val, n, n.left, n.right, hashLeaf(n.key, n.val)) + return hashLeaf(n.key, n.val) + } + fmt.Printf("%*s%d %p %p %p %v\n", depth*2, "", n.bit(), n, n.left, n.right, n.ihash) + h := hashInner(nbit, n.left.check(depth+1, nbit), n.right.check(depth+1, nbit)) + if h != n.ihash && !n.dirty { + fmt.Printf("%*shave %v want %v\n", depth*2, "", n.ihash, h) + panic("bad hash") + } + return h +} diff --git a/mpt/mptload.go b/mpt/mptload.go new file mode 100644 index 0000000..863e291 --- /dev/null +++ b/mpt/mptload.go @@ -0,0 +1,64 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +//go:build ignore + +// Mptload loads a list of keys into a database. +package main + +import ( + "bytes" + "crypto/sha256" + "flag" + "fmt" + "log" + "os" + + "filippo.io/torchwood/mpt" +) + +func usage() { + fmt.Fprintf(os.Stderr, "usage: mptload db1 db2 keys.txt\n") + os.Exit(2) +} + +func main() { + log.SetPrefix("mptload: ") + flag.Usage = usage + flag.Parse() + if flag.NArg() != 3 { + usage() + } + + file1, file2, keys := flag.Arg(0), flag.Arg(1), flag.Arg(2) + tree, err := mpt.Create(file1, file2) + if err != nil { + log.Fatal(err) + } + + data, err := os.ReadFile(keys) + if err != nil { + log.Fatal(err) + } + n := 0 + for line := range bytes.Lines(data) { + h := sha256.Sum256(line) + if err := tree.Set(h, h); err != nil { + log.Fatal(err) + } + n++ + if n%1000000 == 0 { + log.Printf("stored %d", n) + } + } + log.Print("snap") + if _, err := tree.Snap(); err != nil { + log.Fatal(err) + } + log.Print("sync") + if err := tree.Sync(); err != nil { + log.Fatal(err) + } + log.Print("done") +} diff --git a/mpt/tree.go b/mpt/tree.go new file mode 100644 index 0000000..a0d9c4d --- /dev/null +++ b/mpt/tree.go @@ -0,0 +1,374 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package mpt implements a Merkle Patricia Tree. +package mpt + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "iter" + "math/bits" +) + +// A Tree is a Merkle Patricia Tree implementation. +type Tree interface { + // Set adds the given key-value pair to the tree. + // If there is already an entry for the given key, + // then val replaces the old value. + // + // Set is a mutating operation and must not be called + // concurrently with any other Tree method calls + // (including other calls to Set). + Set(key Key, val Val) error + + // Predict returns the hash of the tree that would result from + // applying the given changes (sorted by key) to the tree, + // without modifying the tree. + // + // It is an error to call Predict if Set has been called without + // a subsequent call to Snap: in that case, the caller does not + // know what the current hash is. + Predict(changes []KeyVal) (Hash, error) + + // Snap sets the tree's version number and returns the current tree snapshot. + // + // Snap is a mutating operation and must not be called + // concurrently with any other Tree method calls + // (including other calls to Snap). + // + // As a special case, if version is negative, Snap does not + // set the version. + Snap(version int64) (Snapshot, error) + + // Prove looks up key in the tree and returns a proof + // either of key's value or that key is not present. + // Use [Verify] to retrieve the lookup result. + // + // Prove is a read-only operation and can be called + // concurrently with other calls to Prove, but not other + // calls to Set or Snap. + // + // It is an error to call Prove if Set has been called without + // a subsequent call to Snap: in that case, the caller does not + // know what the root hash is, so the proof will be unverifiable. + Prove(key Key) (Proof, error) + + // Sync flushes all changes from past Set and Snap calls to + // the underlying files and then calls the files' Sync methods + // to flush the changes to disk. (If the files are *os.File files, + // Sync calls fsync(2).) + // + // Even in the absence of calls to Sync, a Tree provides the + // guarantee that on recovery from a crash, it can identify the + // latest snapshot whose Set calls are fully included in the tree. + // A client can call Version() to find the stored tree's version V + // and a boolean indicating whether any later Set calls may also + // be reflected in the tree. When a recovered version is inexact, + // some Set calls made after that version may be present and + // others may not be, no matter the order in which the Set calls were made. + Sync() error + + // Version returns the version number of the tree's last complete snapshot. + // All Set calls made prior to Snap(version) are guaranteed to be + // recorded in the tree. However, if exact is false, then the tree may + // include the effect of Set calls made after that snapshot. + // In that case, to bring the tree into a consistent state, the client is + // expected to replay all Set calls up to the next version. + Version() (version int64, exact bool) + + // Close calls Sync and then closes the underlying files. + Close() error +} + +// ErrModifiedTree indicates that Prove was called after a Set without a Snap. +var ErrModifiedTree = errors.New("tree modified without snapshot") + +// A Key is a key used by a Tree. +// It is usually a cryptographic hash of the actual key data. +type Key [32]byte + +// keyBits is the number of bits in a Key. +const keyBits = len(Key{}) * 8 + +func (k Key) String() string { + return hex.EncodeToString(k[:]) +} + +// bit returns the n'th bit of the key. +func (k Key) bit(n int) int { + return (int(k[n>>3]) >> (7 - n&7)) & 1 +} + +// overlap returns the number of leading bits p and q have in common. +func (p Key) overlap(q Key) int { + for i := range p { + pf := p[i] + qf := q[i] + if pf != qf { + return i*8 + bits.LeadingZeros8(pf^qf) + + } + } + return 256 +} + +// Compare returns the result of comparing two keys. +func (k Key) Compare(q Key) int { + return bytes.Compare(k[:], q[:]) +} + +// Value is the old name for Val. +// Run “go fix” to update client code to use Val instead of Value. +// +//go:fix inline +type Value = Val + +// A Val is a value stored in a Tree. +// It is usually a cryptographic hash of the actual value data. +type Val [32]byte + +func (v Val) String() string { + return hex.EncodeToString(v[:]) +} + +// KeyVal is a key-value pair. +type KeyVal struct { + Key Key + Val Val +} + +// Compare returns the result of comparing keys kv.Key and other.Key. +// It ignores the Val fields. +func (kv KeyVal) Compare(other KeyVal) int { + return kv.Key.Compare(other.Key) +} + +// A keyPrefix is a prefix of a key, identifying a specific node. +type keyPrefix struct { + // bits is the prefix length in bits (0..256, inclusive). + bits int + + // full is the key prefix bytes, zero-padded on the right. + full Key +} + +func (p keyPrefix) String() string { + return fmt.Sprintf("%x/%d", p.full[:(p.bits+7)/8], p.bits) +} + +// overlap returns the number of leading bits p and q have in common. +func (p keyPrefix) overlap(q keyPrefix) int { + return min(p.bits, q.bits, p.full.overlap(q.full)) +} + +func (p keyPrefix) truncate(bits int) keyPrefix { + p.bits = bits + clear(p.full[(bits+7)/8:]) + if n := bits & 7; n != 0 { + p.full[bits/8] &= 0xFF << (8 - n) + } + return p +} + +func (p keyPrefix) compare(q keyPrefix) int { + return bytes.Compare(p.full[:], q.full[:]) +} + +func prefix(key Key, bits int) keyPrefix { + p := keyPrefix{bits: bits, full: key} + return p.truncate(bits) +} + +// A node represents the metadata for a single node. +type node struct { + key keyPrefix + hash Hash +} + +func (x node) merge(y node) node { + b := x.key.overlap(y.key) + return node{x.key.truncate(b), hashInner(b, x.hash, y.hash)} +} + +// A Snapshot is a cryptographic snapshot of a Tree at a point in time. +// It is expected that every snapshot is recorded in a transparent log. +// +// The snapshot epoch is a sequence number identifying a specific snapshot. +// An empty Tree has epoch 0, and then the epoch is incremented each +// time a new snapshot is created (by calling [Tree.Snap] after new records +// are added). +// +// The snapshot hash is a cryptographic hash of the entire tree content. +type Snapshot struct { + Version int64 + Hash Hash +} + +// A Hash is a Merkle hash of a node. +type Hash [32]byte + +func (h Hash) String() string { + return hex.EncodeToString(h[:]) +} + +// TreeHash computes the snapshot hash of a tree consisting of +// the sequence of key-value items. +// +// The sequence must be sorted by increasing +// key value (such as by [Key.Compare] or [KeyVal.Compare]), +// and a key cannot appear multiple times in the list. +// TreeHash panics if the sequence is not sorted or a key appears twice. +// +// Use [slices.Values] to apply TreeHash to a slice of KeyVal. +func TreeHash(seq iter.Seq[KeyVal]) Hash { + var s []node + for kv := range seq { + s = reduce(append(s, node{prefix(kv.Key, keyBits), hashLeaf(kv.Key, kv.Val)})) + } + return hashStack(s) +} + +// A Proof is a proof of the result of looking up a target key in a +// specific snapshot of a Tree. +type Proof []byte + +// Proof Format +// +// Proofs start with "mptproof", followed by a one-byte tag that determines +// the format of the additional data. The tags are: +// +// - 0: proof of empty tree; no data +// - 1: proof key is in tree; data is value and path +// - 2: proof key in not in tree; data is alt key, value, and path +// +// The proof of an empty tree carries no data; to verify the proof is to check that the +// tree snapshot is the empty tree hash. +// +// The proof of a key being in the tree is the key's value followed by the +// path from that key-value pair up to the tree root. +// For each node along the path, the data contains a one-byte overlap count +// (the number of bits shared by the left and right children of the node) +// and the 32-byte hash of the sibling not on the path. +// Verifying the proof requires computing the leaf hash corresponding to key-value +// and then combining that leaf hash with the overlap counts and sibling hashes, +// eventually producing a root tree hash that must match the tree snapshot. +// +// The proof of a key not being in the tree is an alternate key-value pair +// followed by the path from that key-value pair up to the tree root. +// Verifying the proof requires checking that the alt-key is not equal to the +// target key, then recomputing the tree hash from alt-key-value and path. +// During the recomputation, the verifier must check that for every overlap count +// in the path, the target key and the alt-key agree at that bit position, +// verifying that a search for the target would find the alt-key instead. +const ( + proofMagic = "mptproof" + proofEmpty = proofMagic + "\x00" + proofConfirm = proofMagic + "\x01" + proofDeny = proofMagic + "\x02" +) + +var ( + // ErrMalformedProof indicates that a proof is not formatted correctly. + ErrMalformedProof = errors.New("malformed mpt proof") + + // ErrMismatchedProof indicates that a proof does not match + // the snapshot and key passed to Verify. + ErrMismatchedProof = errors.New("mismatched mpt proof") +) + +// Verify verifies that p is a valid proof of a lookup for key in snap, +// returning the proved lookup result (val, ok). +// If the proof is not valid for key in snap, Verify returns a non-nil error. +func Verify(snap Snapshot, key Key, proof Proof) (val Val, ok bool, err error) { + if string(proof) == proofEmpty { + if snap.Hash == emptyTreeHash() { + return Val{}, false, nil + } + return Val{}, false, ErrMismatchedProof + } + + var data []byte + var pkey Key + if data, ok = bytes.CutPrefix(proof, []byte(proofConfirm)); ok && len(data) >= 32 { + pkey = key + val, data = Val(data[:32]), data[32:] + } else if data, ok = bytes.CutPrefix(proof, []byte(proofDeny)); ok && len(data) >= 64 { + pkey, val, data = Key(data[:32]), Val(data[32:64]), data[64:] + if pkey == key { + return Val{}, false, ErrMalformedProof + } + } + h := hashLeaf(pkey, val) + b := 256 + for len(data) >= 1+32 && int(data[0]) < b { + var sib Hash + b, sib, data = int(data[0]), Hash(data[1:1+32]), data[1+32:] + if key.bit(b) != pkey.bit(b) { + return Val{}, false, ErrMalformedProof + } + if key.bit(b) == 0 { + h = hashInner(b, h, sib) + } else { + h = hashInner(b, sib, h) + } + } + if len(data) != 0 || h != snap.Hash { + return Val{}, false, ErrMalformedProof + } + if pkey == key { + return val, true, nil + } + return Val{}, false, nil +} + +// emptyTreeHash returns the parent hash for a root no child nodes. +func emptyTreeHash() Hash { + h := sha256.Sum256(nil) + return h +} + +// hashLeaf returns the hash of a leaf with a given key and value. +func hashLeaf(key Key, val Val) Hash { + var kv [64]byte + copy(kv[:32], key[:]) + copy(kv[32:64], val[:]) + h := sha256.Sum256(kv[:]) + return h +} + +// hashInner returns the hash of an inner node +// with the given bit position and left and right child hashes. +func hashInner(b int, left, right Hash) Hash { + var enc [65]byte + copy(enc[:32], left[:]) + copy(enc[32:64], right[:]) + enc[64] = byte(b) + h := sha256.Sum256(enc[:]) + if right == (Hash{}) { + panic("zero") + } + return h +} + +func reduce(s []node) []node { + for len(s) >= 3 && s[len(s)-3].key.overlap(s[len(s)-2].key) > s[len(s)-2].key.overlap(s[len(s)-1].key) { + m := s[len(s)-3].merge(s[len(s)-2]) + s = append(s[:len(s)-3], m, s[len(s)-1]) + } + return s +} + +func hashStack(s []node) Hash { + if len(s) == 0 { + return emptyTreeHash() + } + for len(s) >= 2 { + s = append(s[:len(s)-2], s[len(s)-2].merge(s[len(s)-1])) + } + return s[0].hash +} diff --git a/mpt/tree_test.go b/mpt/tree_test.go new file mode 100644 index 0000000..647f906 --- /dev/null +++ b/mpt/tree_test.go @@ -0,0 +1,450 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package mpt + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "math/rand" + "runtime/debug" + "slices" + "strings" + "testing" +) + +var goldenTrees = []struct { + keys []Key + hash Hash +}{ + { + []Key{}, + // sha256 /dev/null + sha(), + }, + { + []Key{h("0...0")}, + sha(h("00...0"), h("420...0")), + }, + { + []Key{h("80...0")}, + sha(h("80...0"), h("420...0")), + }, + { + []Key{h("0...0"), h("80...0")}, + sha( + sha(h("00...0"), h("420...0")), + sha(h("80...0"), h("420...01")), + "\x00", + ), + }, + { + []Key{h("0...0"), h("0010...0")}, + sha( + sha(h("0...0"), h("420...0")), + sha(h("0010...0"), h("420...01")), + "\x0b", + ), + }, + { + []Key{h("0...0"), h("0010...0"), h("80...0")}, + sha( + sha( + sha(h("0...0"), h("420...0")), + sha(h("0010...0"), h("420...01")), + "\x0b", + ), + sha(h("80...0"), h("420...02")), + "\x00", + ), + }, +} + +var missing = []Key{ + h("02...2"), + h("22...2"), + h("42...2"), + h("62...2"), + h("82...2"), + h("a2...2"), + h("c2...2"), + h("e2...2"), + h("f2...2"), +} + +func testImpls(t *testing.T, run func(*testing.T, func(*testing.T) *testTree)) { + t.Run("impl=mem", func(t *testing.T) { run(t, testMemTree) }) + t.Run("impl=disk", func(t *testing.T) { run(t, testDiskTree) }) +} + +func TestGoldenTrees(t *testing.T) { + testImpls(t, func(t *testing.T, newTree func(*testing.T) *testTree) { + for i, tree := range goldenTrees { + t.Run(fmt.Sprint(i), func(t *testing.T) { + tt := newTree(t) + defer tt.tree.Close() + e := int64(1) + if len(tree.keys) == 0 { + e = 0 + } + for i, k := range tree.keys { + tt.set(k, v(i)) + } + tt.snap(e, tree.hash) + for i, k := range tree.keys { + tt.get(k, v(i), true) + } + for _, k := range missing { + tt.get(k, Val{}, false) + } + }) + } + }) +} + +func TestAllTrees(t *testing.T) { + testImpls(t, func(t *testing.T, newTree func(*testing.T) *testTree) { + for _, keys := range []string{"hi", "lo"} { + t.Run(keys, func(t *testing.T) { + const B = 3 + const N = 1 << B + k := func(i int) Key { + var k Key + if keys == "hi" { + k[0] = byte(i) << (8 - B) + } else { + k[len(k)-1] = byte(i) + } + return k + } + for leaves := range 1 << N { + tt := newTree(t) + var kvs []KeyVal + for i := range N { + if leaves&(1<