Line data Source code
1 : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package manifest
6 :
7 : // The Annotator type defined below is used by other packages to lazily
8 : // compute a value over a B-Tree. Each node of the B-Tree stores one
9 : // `annotation` per annotator, containing the result of the computation over
10 : // the node's subtree.
11 : //
12 : // An annotation is marked as valid if it's current with the current subtree
13 : // state. Annotations are marked as invalid whenever a node will be mutated
14 : // (in mut). Annotators may also return `false` from `Accumulate` to signal
15 : // that a computation for a file is not stable and may change in the future.
16 : // Annotations that include these unstable values are also marked as invalid
17 : // on the node, ensuring that future queries for the annotation will recompute
18 : // the value.
19 :
20 : // An Annotator defines a computation over a level's FileMetadata. If the
21 : // computation is stable and uses inputs that are fixed for the lifetime of
22 : // a FileMetadata, the LevelMetadata's internal data structures are annotated
23 : // with the intermediary computations. This allows the computation to be
24 : // computed incrementally as edits are applied to a level.
25 : type Annotator[T any] struct {
26 : Aggregator AnnotationAggregator[T]
27 : }
28 :
29 : // An AnnotationAggregator defines how an annotation should be accumulated
30 : // from a single FileMetadata and merged with other annotated values.
31 : type AnnotationAggregator[T any] interface {
32 : // Zero returns the zero value of an annotation. This value is returned
33 : // when a LevelMetadata is empty. The dst argument, if non-nil, is an
34 : // obsolete value previously returned by this Annotator and may be
35 : // overwritten and reused to avoid a memory allocation.
36 : Zero(dst *T) *T
37 :
38 : // Accumulate computes the annotation for a single file in a level's
39 : // metadata. It merges the file's value into dst and returns a bool flag
40 : // indicating whether or not the value is stable and okay to cache as an
41 : // annotation. If the file's value may change over the life of the file,
42 : // the annotator must return false.
43 : //
44 : // Implementations may modify dst and return it to avoid an allocation.
45 : Accumulate(f *FileMetadata, dst *T) (v *T, cacheOK bool)
46 :
47 : // Merge combines two values src and dst, returning the result.
48 : // Implementations may modify dst and return it to avoid an allocation.
49 : Merge(src *T, dst *T) *T
50 : }
51 :
52 : type annotation struct {
53 : // annotator is a pointer to the Annotator that computed this annotation.
54 : // NB: This is untyped to allow AnnotationAggregator to use Go generics,
55 : // since annotations are stored in a slice on each node and a single
56 : // slice cannot contain elements with different type parameters.
57 : annotator interface{}
58 : // v is contains the annotation value, the output of either
59 : // AnnotationAggregator.Accumulate or AnnotationAggregator.Merge.
60 : // NB: This is untyped for the same reason as annotator above.
61 : v interface{}
62 : // valid indicates whether future reads of the annotation may use the
63 : // value as-is. If false, v will be zeroed and recalculated.
64 : valid bool
65 : }
66 :
67 : // findAnnotation finds this Annotator's annotation on a node, creating
68 : // one if it doesn't already exist.
69 1 : func (a *Annotator[T]) findAnnotation(n *node) *annotation {
70 1 : for i := range n.annot {
71 1 : if n.annot[i].annotator == a {
72 1 : return &n.annot[i]
73 1 : }
74 : }
75 :
76 : // This node has never been annotated by a. Create a new annotation.
77 1 : n.annot = append(n.annot, annotation{
78 1 : annotator: a,
79 1 : v: a.Aggregator.Zero(nil),
80 1 : })
81 1 : return &n.annot[len(n.annot)-1]
82 : }
83 :
84 : // nodeAnnotation computes this annotator's annotation of this node across all
85 : // files in the node's subtree. The second return value indicates whether the
86 : // annotation is stable and thus cacheable.
87 1 : func (a *Annotator[T]) nodeAnnotation(n *node) (_ *T, cacheOK bool) {
88 1 : annot := a.findAnnotation(n)
89 1 : t := annot.v.(*T)
90 1 : // If the annotation is already marked as valid, we can return it without
91 1 : // recomputing anything.
92 1 : if annot.valid {
93 1 : return t, true
94 1 : }
95 :
96 1 : t = a.Aggregator.Zero(t)
97 1 : valid := true
98 1 :
99 1 : for i := int16(0); i <= n.count; i++ {
100 1 : if !n.leaf {
101 1 : v, ok := a.nodeAnnotation(n.children[i])
102 1 : t = a.Aggregator.Merge(v, t)
103 1 : valid = valid && ok
104 1 : }
105 :
106 1 : if i < n.count {
107 1 : var ok bool
108 1 : t, ok = a.Aggregator.Accumulate(n.items[i], t)
109 1 : valid = valid && ok
110 1 : }
111 : }
112 :
113 1 : annot.v = t
114 1 : annot.valid = valid
115 1 :
116 1 : return t, annot.valid
117 : }
118 :
119 : // InvalidateAnnotation removes any existing cached annotations from this
120 : // annotator from a node's subtree.
121 1 : func (a *Annotator[T]) invalidateNodeAnnotation(n *node) {
122 1 : annot := a.findAnnotation(n)
123 1 : annot.valid = false
124 1 : if !n.leaf {
125 0 : for i := int16(0); i <= n.count; i++ {
126 0 : a.invalidateNodeAnnotation(n.children[i])
127 0 : }
128 : }
129 : }
130 :
131 : // LevelAnnotation calculates the annotation defined by this Annotator for all
132 : // files in the given LevelMetadata. A pointer to the Annotator is used as the
133 : // key for pre-calculated values, so the same Annotator must be used to avoid
134 : // duplicate computation. Annotation must not be called concurrently, and in
135 : // practice this is achieved by requiring callers to hold DB.mu.
136 1 : func (a *Annotator[T]) LevelAnnotation(lm LevelMetadata) *T {
137 1 : if lm.Empty() {
138 1 : return a.Aggregator.Zero(nil)
139 1 : }
140 :
141 1 : v, _ := a.nodeAnnotation(lm.tree.root)
142 1 : return v
143 : }
144 :
145 : // LevelAnnotation calculates the annotation defined by this Annotator for all
146 : // files across the given levels. A pointer to the Annotator is used as the
147 : // key for pre-calculated values, so the same Annotator must be used to avoid
148 : // duplicate computation. Annotation must not be called concurrently, and in
149 : // practice this is achieved by requiring callers to hold DB.mu.
150 1 : func (a *Annotator[T]) MultiLevelAnnotation(lms []LevelMetadata) *T {
151 1 : aggregated := a.Aggregator.Zero(nil)
152 1 : for l := 0; l < len(lms); l++ {
153 1 : if !lms[l].Empty() {
154 1 : v := a.LevelAnnotation(lms[l])
155 1 : aggregated = a.Aggregator.Merge(v, aggregated)
156 1 : }
157 : }
158 1 : return aggregated
159 : }
160 :
161 : // InvalidateAnnotation clears any cached annotations defined by Annotator. A
162 : // pointer to the Annotator is used as the key for pre-calculated values, so
163 : // the same Annotator must be used to clear the appropriate cached annotation.
164 : // InvalidateAnnotation must not be called concurrently, and in practice this
165 : // is achieved by requiring callers to hold DB.mu.
166 1 : func (a *Annotator[T]) InvalidateLevelAnnotation(lm LevelMetadata) {
167 1 : if lm.Empty() {
168 0 : return
169 0 : }
170 1 : a.invalidateNodeAnnotation(lm.tree.root)
171 : }
172 :
173 : // sumAggregator defines an Aggregator which sums together a uint64 value
174 : // across files.
175 : type sumAggregator struct {
176 : accumulateFunc func(f *FileMetadata) (v uint64, cacheOK bool)
177 : }
178 :
179 1 : func (sa sumAggregator) Zero(dst *uint64) *uint64 {
180 1 : if dst == nil {
181 1 : return new(uint64)
182 1 : }
183 1 : *dst = 0
184 1 : return dst
185 : }
186 :
187 1 : func (sa sumAggregator) Accumulate(f *FileMetadata, dst *uint64) (v *uint64, cacheOK bool) {
188 1 : accumulated, ok := sa.accumulateFunc(f)
189 1 : *dst += accumulated
190 1 : return dst, ok
191 1 : }
192 :
193 1 : func (sa sumAggregator) Merge(src *uint64, dst *uint64) *uint64 {
194 1 : *dst += *src
195 1 : return dst
196 1 : }
197 :
198 : // SumAnnotator takes a function that computes a uint64 value from a single
199 : // FileMetadata and returns an Annotator that sums together the values across
200 : // files.
201 1 : func SumAnnotator(accumulate func(f *FileMetadata) (v uint64, cacheOK bool)) *Annotator[uint64] {
202 1 : return &Annotator[uint64]{
203 1 : Aggregator: sumAggregator{
204 1 : accumulateFunc: accumulate,
205 1 : },
206 1 : }
207 1 : }
208 :
209 : // PickFileAggregator implements the AnnotationAggregator interface. It defines
210 : // an aggregator that picks a single file from a set of eligible files.
211 : type PickFileAggregator struct {
212 : // Filter takes a FileMetadata and returns whether it is eligible to be
213 : // picked by this PickFileAggregator. The second return value indicates
214 : // whether this eligibility is stable and thus cacheable.
215 : Filter func(f *FileMetadata) (eligible bool, cacheOK bool)
216 : // Compare compares two instances of FileMetadata and returns true if
217 : // the first one should be picked over the second one. It may assume
218 : // that both arguments are non-nil.
219 : Compare func(f1 *FileMetadata, f2 *FileMetadata) bool
220 : }
221 :
222 : // Zero implements AnnotationAggregator.Zero, returning nil as the zero value.
223 1 : func (fa PickFileAggregator) Zero(dst *FileMetadata) *FileMetadata {
224 1 : return nil
225 1 : }
226 :
227 1 : func (fa PickFileAggregator) mergePickedFiles(src *FileMetadata, dst *FileMetadata) *FileMetadata {
228 1 : switch {
229 1 : case src == nil:
230 1 : return dst
231 1 : case dst == nil:
232 1 : return src
233 1 : case fa.Compare(src, dst):
234 1 : return src
235 1 : default:
236 1 : return dst
237 : }
238 : }
239 :
240 : // Accumulate implements AnnotationAggregator.Accumulate, accumulating a single
241 : // file as long as it is eligible to be picked.
242 : func (fa PickFileAggregator) Accumulate(
243 : f *FileMetadata, dst *FileMetadata,
244 1 : ) (v *FileMetadata, cacheOK bool) {
245 1 : eligible, ok := fa.Filter(f)
246 1 : if eligible {
247 1 : return fa.mergePickedFiles(f, dst), ok
248 1 : }
249 1 : return dst, ok
250 : }
251 :
252 : // Merge implements AnnotationAggregator.Merge by picking a single file based
253 : // on the output of PickFileAggregator.Compare.
254 1 : func (fa PickFileAggregator) Merge(src *FileMetadata, dst *FileMetadata) *FileMetadata {
255 1 : return fa.mergePickedFiles(src, dst)
256 1 : }
|