github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/metamorphic/retryable.go (about)

     1  // Copyright 2020 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 metamorphic
     6  
     7  import (
     8  	"github.com/cockroachdb/errors"
     9  	"github.com/cockroachdb/pebble"
    10  	"github.com/cockroachdb/pebble/vfs/errorfs"
    11  )
    12  
    13  // withRetries executes fn, retrying it whenever an errorfs.ErrInjected error
    14  // is returned.  It returns the first nil or non-errorfs.ErrInjected error
    15  // returned by fn.
    16  func withRetries(fn func() error) error {
    17  	for {
    18  		if err := fn(); !errors.Is(err, errorfs.ErrInjected) {
    19  			return err
    20  		}
    21  	}
    22  }
    23  
    24  // retryableIter holds an iterator and the state necessary to reset it to its
    25  // state after the last successful operation. This allows us to retry failed
    26  // iterator operations by running them again on a non-error iterator with the
    27  // same pre-operation state.
    28  type retryableIter struct {
    29  	iter    *pebble.Iterator
    30  	lastKey []byte
    31  }
    32  
    33  func (i *retryableIter) needRetry() bool {
    34  	return errors.Is(i.iter.Error(), errorfs.ErrInjected)
    35  }
    36  
    37  func (i *retryableIter) withRetry(fn func()) {
    38  	for {
    39  		fn()
    40  		if !i.needRetry() {
    41  			break
    42  		}
    43  		for i.needRetry() {
    44  			i.iter.SeekGE(i.lastKey)
    45  		}
    46  	}
    47  
    48  	i.lastKey = i.lastKey[:0]
    49  	if i.iter.Valid() {
    50  		i.lastKey = append(i.lastKey, i.iter.Key()...)
    51  	}
    52  }
    53  
    54  func (i *retryableIter) Close() error {
    55  	return i.iter.Close()
    56  }
    57  
    58  func (i *retryableIter) Error() error {
    59  	return i.iter.Error()
    60  }
    61  
    62  func (i *retryableIter) First() bool {
    63  	var valid bool
    64  	i.withRetry(func() {
    65  		valid = i.iter.First()
    66  	})
    67  	return valid
    68  }
    69  
    70  func (i *retryableIter) Key() []byte {
    71  	return i.iter.Key()
    72  }
    73  
    74  func (i *retryableIter) RangeKeyChanged() bool {
    75  	return i.iter.RangeKeyChanged()
    76  }
    77  
    78  func (i *retryableIter) HasPointAndRange() (bool, bool) {
    79  	return i.iter.HasPointAndRange()
    80  }
    81  
    82  func (i *retryableIter) RangeBounds() ([]byte, []byte) {
    83  	return i.iter.RangeBounds()
    84  }
    85  
    86  func (i *retryableIter) RangeKeys() []pebble.RangeKeyData {
    87  	return i.iter.RangeKeys()
    88  }
    89  
    90  func (i *retryableIter) Last() bool {
    91  	var valid bool
    92  	i.withRetry(func() { valid = i.iter.Last() })
    93  	return valid
    94  }
    95  
    96  func (i *retryableIter) Next() bool {
    97  	var valid bool
    98  	i.withRetry(func() {
    99  		valid = i.iter.Next()
   100  	})
   101  	return valid
   102  }
   103  
   104  func (i *retryableIter) NextWithLimit(limit []byte) pebble.IterValidityState {
   105  	var validity pebble.IterValidityState
   106  	i.withRetry(func() {
   107  		validity = i.iter.NextWithLimit(limit)
   108  	})
   109  	return validity
   110  
   111  }
   112  
   113  func (i *retryableIter) NextPrefix() bool {
   114  	var valid bool
   115  	i.withRetry(func() {
   116  		valid = i.iter.NextPrefix()
   117  	})
   118  	return valid
   119  }
   120  
   121  func (i *retryableIter) Prev() bool {
   122  	var valid bool
   123  	i.withRetry(func() {
   124  		valid = i.iter.Prev()
   125  	})
   126  	return valid
   127  }
   128  
   129  func (i *retryableIter) PrevWithLimit(limit []byte) pebble.IterValidityState {
   130  	var validity pebble.IterValidityState
   131  	i.withRetry(func() {
   132  		validity = i.iter.PrevWithLimit(limit)
   133  	})
   134  	return validity
   135  }
   136  
   137  func (i *retryableIter) SeekGE(key []byte) bool {
   138  	var valid bool
   139  	i.withRetry(func() { valid = i.iter.SeekGE(key) })
   140  	return valid
   141  }
   142  
   143  func (i *retryableIter) SeekGEWithLimit(key []byte, limit []byte) pebble.IterValidityState {
   144  	var validity pebble.IterValidityState
   145  	i.withRetry(func() { validity = i.iter.SeekGEWithLimit(key, limit) })
   146  	return validity
   147  }
   148  
   149  func (i *retryableIter) SeekLT(key []byte) bool {
   150  	var valid bool
   151  	i.withRetry(func() { valid = i.iter.SeekLT(key) })
   152  	return valid
   153  }
   154  
   155  func (i *retryableIter) SeekLTWithLimit(key []byte, limit []byte) pebble.IterValidityState {
   156  	var validity pebble.IterValidityState
   157  	i.withRetry(func() { validity = i.iter.SeekLTWithLimit(key, limit) })
   158  	return validity
   159  }
   160  
   161  func (i *retryableIter) SeekPrefixGE(key []byte) bool {
   162  	var valid bool
   163  	i.withRetry(func() { valid = i.iter.SeekPrefixGE(key) })
   164  	return valid
   165  }
   166  
   167  func (i *retryableIter) SetBounds(lower, upper []byte) {
   168  	i.iter.SetBounds(lower, upper)
   169  }
   170  
   171  func (i *retryableIter) SetOptions(opts *pebble.IterOptions) {
   172  	i.iter.SetOptions(opts)
   173  }
   174  
   175  func (i *retryableIter) Valid() bool {
   176  	return i.iter.Valid()
   177  }
   178  
   179  func (i *retryableIter) Value() []byte {
   180  	return i.iter.Value()
   181  }