go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/featureBreaker/rds.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 featureBreaker
    16  
    17  import (
    18  	"context"
    19  
    20  	ds "go.chromium.org/luci/gae/service/datastore"
    21  )
    22  
    23  // DatastoreFeatures is a list of datastore features that can be "broken".
    24  var DatastoreFeatures = []string{
    25  	"AllocateIDs",
    26  	"DecodeCursor",
    27  	"Run",
    28  	"Count",
    29  	"BeginTransaction",
    30  	"CommitTransaction",
    31  	"DeleteMulti",
    32  	"GetMulti",
    33  	"PutMulti",
    34  }
    35  
    36  type dsState struct {
    37  	*state
    38  
    39  	c   context.Context
    40  	rds ds.RawInterface
    41  }
    42  
    43  func (r *dsState) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
    44  	if len(keys) == 0 {
    45  		return nil
    46  	}
    47  	return r.run(r.c, func() error {
    48  		return r.rds.AllocateIDs(keys, cb)
    49  	})
    50  }
    51  
    52  func (r *dsState) DecodeCursor(s string) (ds.Cursor, error) {
    53  	curs := ds.Cursor(nil)
    54  	err := r.run(r.c, func() (err error) {
    55  		curs, err = r.rds.DecodeCursor(s)
    56  		return
    57  	})
    58  	return curs, err
    59  }
    60  
    61  func (r *dsState) Run(q *ds.FinalizedQuery, cb ds.RawRunCB) error {
    62  	return r.run(r.c, func() error {
    63  		return r.rds.Run(q, cb)
    64  	})
    65  }
    66  
    67  func (r *dsState) Count(q *ds.FinalizedQuery) (int64, error) {
    68  	count := int64(0)
    69  	err := r.run(r.c, func() (err error) {
    70  		count, err = r.rds.Count(q)
    71  		return
    72  	})
    73  	return count, err
    74  }
    75  
    76  func (r *dsState) RunInTransaction(f func(c context.Context) error, opts *ds.TransactionOptions) error {
    77  	// Note: we intentionally don't break RunInTransaction itself, but break
    78  	// BeginTransaction/CommitTransaction separately instead.
    79  	return r.rds.RunInTransaction(func(txnc context.Context) error {
    80  		if err := r.BeginTransaction(txnc); err != nil {
    81  			return err
    82  		}
    83  		if err := f(txnc); err != nil {
    84  			return err
    85  		}
    86  		return r.CommitTransaction(txnc)
    87  	}, opts)
    88  }
    89  
    90  // BeginTransaction is exposed as a "breakable" hook to simulate bad transaction
    91  // initiation.
    92  //
    93  // It not a part of RawDatastore interface, but can nevertheless be overridden
    94  // with BreakFeature to simulate "BeginTransaction" RPC errors.
    95  //
    96  // Unlike RunInTransaction (which is called once), BeginTransaction is called at
    97  // the beginning of each individual transaction retry.
    98  func (r *dsState) BeginTransaction(c context.Context) error {
    99  	return r.run(c, func() error { return nil })
   100  }
   101  
   102  // CommitTransaction is exposed as a "breakable" hook to simulate bad
   103  // transaction commits.
   104  //
   105  // It is not a part of RawDatastore interface, but can nevertheless be
   106  // overridden with BreakFeature to simulate transaction commit errors.
   107  func (r *dsState) CommitTransaction(c context.Context) error {
   108  	return r.run(c, func() error { return nil })
   109  }
   110  
   111  // TODO(iannucci): Allow the user to specify a multierror which will propagate
   112  // to the callback correctly.
   113  
   114  func (r *dsState) DeleteMulti(keys []*ds.Key, cb ds.DeleteMultiCB) error {
   115  	if len(keys) == 0 {
   116  		return nil
   117  	}
   118  	return r.run(r.c, func() error {
   119  		return r.rds.DeleteMulti(keys, cb)
   120  	})
   121  }
   122  
   123  func (r *dsState) GetMulti(keys []*ds.Key, meta ds.MultiMetaGetter, cb ds.GetMultiCB) error {
   124  	if len(keys) == 0 {
   125  		return nil
   126  	}
   127  	return r.run(r.c, func() error {
   128  		return r.rds.GetMulti(keys, meta, cb)
   129  	})
   130  }
   131  
   132  func (r *dsState) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
   133  	if len(keys) == 0 {
   134  		return nil
   135  	}
   136  	return r.run(r.c, func() (err error) {
   137  		return r.rds.PutMulti(keys, vals, cb)
   138  	})
   139  }
   140  
   141  func (r *dsState) WithoutTransaction() context.Context {
   142  	return r.rds.WithoutTransaction()
   143  }
   144  
   145  func (r *dsState) CurrentTransaction() ds.Transaction {
   146  	return r.rds.CurrentTransaction()
   147  }
   148  
   149  func (r *dsState) Constraints() ds.Constraints { return r.rds.Constraints() }
   150  
   151  func (r *dsState) GetTestable() ds.Testable {
   152  	return r.rds.GetTestable()
   153  }
   154  
   155  // FilterRDS installs a featureBreaker datastore filter in the context.
   156  func FilterRDS(c context.Context, defaultError error) (context.Context, FeatureBreaker) {
   157  	state := newState(defaultError)
   158  	return ds.AddRawFilters(c, func(ic context.Context, rawDatastore ds.RawInterface) ds.RawInterface {
   159  		return &dsState{state, ic, rawDatastore}
   160  	}), state
   161  }