go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/impl/memory/memstore_iter.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package memory
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"go.chromium.org/luci/common/data/cmpbin"
    21  )
    22  
    23  type iterDefinition struct {
    24  	// The collection to iterate over
    25  	c memCollection
    26  
    27  	// The prefix to always assert for every row. A nil prefix matches every row.
    28  	prefix []byte
    29  
    30  	// prefixLen is the number of prefix bytes that the caller cares about. It
    31  	// may be <= len(prefix). When doing a multiIterator, this number will be used
    32  	// to determine the amount of suffix to transfer accross iterators. This is
    33  	// used specifically when using builtin indexes to service ancestor queries.
    34  	// The builtin index represents the ancestor key with prefix bytes, but in a
    35  	// multiIterator context, it wants the entire key to be included in the
    36  	// suffix.
    37  	prefixLen int
    38  
    39  	// The start cursor. It's appended to prefix to find the first row.
    40  	start []byte
    41  
    42  	// The end cursor. It's appended to prefix to find the last row (which is not
    43  	// included in the interation result). If this is nil, then there's no end
    44  	// except the natural end of the collection.
    45  	end []byte
    46  }
    47  
    48  func multiIterate(defs []*iterDefinition, cb func(suffix []byte) error) error {
    49  	if len(defs) == 0 {
    50  		return nil
    51  	}
    52  
    53  	ts := make([]*iterator, len(defs))
    54  	prefixLens := make([]int, len(defs))
    55  	for i, def := range defs {
    56  		// bind i so that the defer below doesn't get goofed by the loop variable
    57  		i := i
    58  		ts[i] = def.mkIter()
    59  		prefixLens[i] = def.prefixLen
    60  	}
    61  
    62  	suffix := []byte(nil)
    63  	skip := -1
    64  
    65  MainLoop:
    66  	for {
    67  		for idx, it := range ts {
    68  			if skip >= 0 && skip == idx {
    69  				continue
    70  			}
    71  
    72  			pfxLen := prefixLens[idx]
    73  			it.skip(cmpbin.ConcatBytes(it.def.prefix[:pfxLen], suffix))
    74  			ent := it.next()
    75  			if ent == nil {
    76  				// we hit the end of an iterator, we're now done with the whole
    77  				// query.
    78  				return nil
    79  			}
    80  			sfxRO := ent.key[pfxLen:]
    81  
    82  			if bytes.Compare(sfxRO, suffix) > 0 {
    83  				// this row has a higher suffix than anything we've seen before. Set
    84  				// ourself to be the skip, and resart this loop from the top.
    85  				suffix = append(suffix[:0], sfxRO...)
    86  				skip = idx
    87  				if idx != 0 {
    88  					// no point to restarting on the 0th index
    89  					continue MainLoop
    90  				}
    91  			}
    92  		}
    93  
    94  		if err := cb(suffix); err != nil {
    95  			return err
    96  		}
    97  		suffix = nil
    98  		skip = -1
    99  	}
   100  }
   101  
   102  type iterator struct {
   103  	def  *iterDefinition
   104  	base memIterator
   105  
   106  	start   []byte
   107  	end     []byte
   108  	lastKey []byte
   109  }
   110  
   111  func (def *iterDefinition) mkIter() *iterator {
   112  	if !def.c.IsReadOnly() {
   113  		panic("attempting to make an iterator with r/w collection")
   114  	}
   115  
   116  	it := iterator{
   117  		def: def,
   118  	}
   119  
   120  	// convert the suffixes from the iterDefinition into full rows for the
   121  	// underlying storage.
   122  	it.start = cmpbin.ConcatBytes(def.prefix, def.start)
   123  	if def.end != nil {
   124  		it.end = cmpbin.ConcatBytes(def.prefix, def.end)
   125  	}
   126  	return &it
   127  }
   128  
   129  func (it *iterator) skip(targ []byte) {
   130  	if bytes.Compare(targ, it.start) < 0 {
   131  		targ = it.start
   132  	}
   133  	if it.base == nil || bytes.Compare(targ, it.lastKey) > 0 {
   134  		// If our skip target is >= our last key, then create a new Iterator
   135  		// starting from that target.
   136  		it.base = it.def.c.Iterator(targ)
   137  	}
   138  }
   139  
   140  func (it *iterator) next() *storeEntry {
   141  	if it.base == nil {
   142  		it.skip(nil)
   143  	}
   144  
   145  	ent := it.base.Next()
   146  	switch {
   147  	case ent == nil:
   148  		return nil
   149  
   150  	case !bytes.HasPrefix(ent.key, it.def.prefix):
   151  		// we're no longer in prefix, terminate
   152  		return nil
   153  
   154  	case it.end != nil && bytes.Compare(ent.key, it.end) >= 0:
   155  		// we hit our cap, terminate.
   156  		return nil
   157  
   158  	default:
   159  		it.lastKey = ent.key
   160  		return ent
   161  	}
   162  }