-
Notifications
You must be signed in to change notification settings - Fork 375
/
writer.go
162 lines (147 loc) · 5.71 KB
/
writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// Copyright 2019 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package smt
import (
"context"
"errors"
"fmt"
"github.com/google/trillian/merkle/smt/node"
)
// NodeBatchAccessor reads and writes batches of Merkle tree node hashes. It is
// a batch interface for efficiency reasons, as it is designed to guard tree
// storage / database access. The Writer type operates on a per-shard basis,
// i.e. it calls Get and Set method exactly once for each shard.
type NodeBatchAccessor interface {
// Get returns the hashes of the given nodes, as a map keyed by their IDs.
// The returned hashes may be missing or be nil for empty subtrees.
Get(ctx context.Context, ids []node.ID) (map[node.ID][]byte, error)
// Set applies the given node hash updates.
Set(ctx context.Context, nodes []Node) error
}
// Writer handles sharded writes to a sparse Merkle tree. The tree has two
// levels of shards: the single topmost shard spanning depths from 0 to split,
// and 2^split second-level shards each spanning levels from split to height.
// If the split height is 0 then effectively there is only one "global" shard.
type Writer struct {
h mapHasher
height uint // The height of the tree.
split uint // The height of the top shard.
}
// NewWriter creates a new Writer for the specified tree of the given height,
// with two levels of sharding, where the upper shard is `split` levels high.
func NewWriter(treeID int64, hasher Hasher, height, split uint) *Writer {
if split > height {
panic(fmt.Errorf("NewWriter: split(%d) > height(%d)", split, height))
}
return &Writer{h: bindHasher(hasher, treeID), height: height, split: split}
}
// Split sorts and splits the given list of node hash updates into shards, i.e.
// the subsets belonging to different subtrees. The nodes must belong to the
// same tree level which is equal to the tree height.
func (w *Writer) Split(nodes []Node) ([][]Node, error) {
if err := Prepare(nodes, w.height); err != nil {
return nil, err
}
// TODO(pavelkalinnikov): Try estimating the capacity for this slice.
var shards [][]Node
// The nodes are sorted, so we can split them by prefix.
for begin, i := 0, 0; i < len(nodes); i++ {
pref := nodes[i].ID.Prefix(w.split)
next := i + 1
// Check if this ID ends the shard.
if next == len(nodes) || nodes[next].ID.Prefix(w.split) != pref {
shards = append(shards, nodes[begin:next])
begin = next
}
}
return shards, nil
}
// Write applies the given list of node updates to a single shard, and returns
// the resulting update of the shard root. It uses the given node accessor for
// reading and writing tree nodes.
//
// The typical usage pattern is as follows. For the lower shards, the input is
// the []Node slices returned from the Split method. For the top shard, the
// input is all the Node values from the lower shards Write calls.
//
// In another case, Write can be performed without Split if the shard split
// depth is 0, which effectively means that there is only one "global" shard.
func (w *Writer) Write(ctx context.Context, nodes []Node, acc NodeBatchAccessor) (Node, error) {
if len(nodes) == 0 {
return Node{}, errors.New("nothing to write")
}
depth := nodes[0].ID.BitLen()
top, err := w.shardTop(depth)
if err != nil {
return Node{}, err
}
hs, err := NewHStar3(nodes, w.h.mh.HashChildren, depth, top)
if err != nil {
return Node{}, err
}
hashes, err := acc.Get(ctx, hs.Prepare())
if err != nil {
return Node{}, err
}
sa := w.newAccessor(hashes)
topUpd, err := hs.Update(sa)
if err != nil {
return Node{}, err
} else if ln := len(topUpd); ln != 1 {
return Node{}, fmt.Errorf("writing across %d shards, want 1", ln)
}
if err := acc.Set(ctx, sa.writes); err != nil {
return Node{}, err
}
return topUpd[0], nil
}
// shardTop returns the depth of a shard top based on its bottom depth.
func (w *Writer) shardTop(depth uint) (uint, error) {
switch depth {
case w.height:
return w.split, nil
case w.split:
return 0, nil
}
return 0, fmt.Errorf("unexpected depth %d", depth)
}
// newAccessor returns a NodeAccessor for HStar3 algorithm based on the set of
// preloaded node hashes.
func (w *Writer) newAccessor(nodes map[node.ID][]byte) *shardAccessor {
// For any node that HStar3 reads, it also writes its sibling. Therefore we
// can pre-allocate this many items for the writes slice.
// TODO(pavelkalinnikov): The actual number of written nodes will be slightly
// bigger by at most the number of written leaves. Try allocating precisely.
writes := make([]Node, 0, len(nodes))
return &shardAccessor{w: w, reads: nodes, writes: writes}
}
// shardAccessor provides read and write access to nodes used by HStar3. It
// operates entirely in-memory.
type shardAccessor struct {
w *Writer
reads map[node.ID][]byte
writes []Node
}
// Get returns the hash of the given node from the preloaded map, or a hash of
// an empty subtree at this position if such node is not found.
func (s *shardAccessor) Get(id node.ID) ([]byte, error) {
if hash, ok := s.reads[id]; ok && hash != nil {
return hash, nil
}
return s.w.h.hashEmpty(id), nil
}
// Set adds the given node hash update to the list of writes.
func (s *shardAccessor) Set(id node.ID, hash []byte) {
s.writes = append(s.writes, Node{ID: id, Hash: hash})
}