github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/ekv/fault_injection.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package ekv
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  )
    20  
    21  // InjectionConfig is used for fault injections for KV components.
    22  type InjectionConfig struct {
    23  	sync.RWMutex
    24  	getError    error // ekv.Get() always return this error.
    25  	commitError error // Transaction.Commit() always return this error.
    26  }
    27  
    28  // SetGetError injects an error for all ekv.Get() methods.
    29  func (c *InjectionConfig) SetGetError(err error) {
    30  	c.Lock()
    31  	defer c.Unlock()
    32  
    33  	c.getError = err
    34  }
    35  
    36  // SetCommitError injects an error for all Transaction.Commit() methods.
    37  func (c *InjectionConfig) SetCommitError(err error) {
    38  	c.Lock()
    39  	defer c.Unlock()
    40  	c.commitError = err
    41  }
    42  
    43  // InjectedStore wraps a CausetStorage with injections.
    44  type InjectedStore struct {
    45  	CausetStorage
    46  	cfg *InjectionConfig
    47  }
    48  
    49  // NewInjectedStore creates a InjectedStore with config.
    50  func NewInjectedStore(causetstore CausetStorage, cfg *InjectionConfig) CausetStorage {
    51  	return &InjectedStore{
    52  		CausetStorage: causetstore,
    53  		cfg:           cfg,
    54  	}
    55  }
    56  
    57  // Begin creates an injected Transaction.
    58  func (s *InjectedStore) Begin() (Transaction, error) {
    59  	txn, err := s.CausetStorage.Begin()
    60  	return &InjectedTransaction{
    61  		Transaction: txn,
    62  		cfg:         s.cfg,
    63  	}, err
    64  }
    65  
    66  // BeginWithStartTS creates an injected Transaction with startTS.
    67  func (s *InjectedStore) BeginWithStartTS(startTS uint64) (Transaction, error) {
    68  	txn, err := s.CausetStorage.BeginWithStartTS(startTS)
    69  	return &InjectedTransaction{
    70  		Transaction: txn,
    71  		cfg:         s.cfg,
    72  	}, err
    73  }
    74  
    75  // GetSnapshot creates an injected Snapshot.
    76  func (s *InjectedStore) GetSnapshot(ver Version) (Snapshot, error) {
    77  	snapshot, err := s.CausetStorage.GetSnapshot(ver)
    78  	return &InjectedSnapshot{
    79  		Snapshot: snapshot,
    80  		cfg:      s.cfg,
    81  	}, err
    82  }
    83  
    84  // InjectedTransaction wraps a Transaction with injections.
    85  type InjectedTransaction struct {
    86  	Transaction
    87  	cfg *InjectionConfig
    88  }
    89  
    90  // Get returns an error if cfg.getError is set.
    91  func (t *InjectedTransaction) Get(ctx context.Context, k Key) ([]byte, error) {
    92  	t.cfg.RLock()
    93  	defer t.cfg.RUnlock()
    94  	if t.cfg.getError != nil {
    95  		return nil, t.cfg.getError
    96  	}
    97  	return t.Transaction.Get(ctx, k)
    98  }
    99  
   100  // BatchGet returns an error if cfg.getError is set.
   101  func (t *InjectedTransaction) BatchGet(ctx context.Context, keys []Key) (map[string][]byte, error) {
   102  	t.cfg.RLock()
   103  	defer t.cfg.RUnlock()
   104  	if t.cfg.getError != nil {
   105  		return nil, t.cfg.getError
   106  	}
   107  	return t.Transaction.BatchGet(ctx, keys)
   108  }
   109  
   110  // Commit returns an error if cfg.commitError is set.
   111  func (t *InjectedTransaction) Commit(ctx context.Context) error {
   112  	t.cfg.RLock()
   113  	defer t.cfg.RUnlock()
   114  	if t.cfg.commitError != nil {
   115  		return t.cfg.commitError
   116  	}
   117  	return t.Transaction.Commit(ctx)
   118  }
   119  
   120  // InjectedSnapshot wraps a Snapshot with injections.
   121  type InjectedSnapshot struct {
   122  	Snapshot
   123  	cfg *InjectionConfig
   124  }
   125  
   126  // Get returns an error if cfg.getError is set.
   127  func (t *InjectedSnapshot) Get(ctx context.Context, k Key) ([]byte, error) {
   128  	t.cfg.RLock()
   129  	defer t.cfg.RUnlock()
   130  	if t.cfg.getError != nil {
   131  		return nil, t.cfg.getError
   132  	}
   133  	return t.Snapshot.Get(ctx, k)
   134  }
   135  
   136  // BatchGet returns an error if cfg.getError is set.
   137  func (t *InjectedSnapshot) BatchGet(ctx context.Context, keys []Key) (map[string][]byte, error) {
   138  	t.cfg.RLock()
   139  	defer t.cfg.RUnlock()
   140  	if t.cfg.getError != nil {
   141  		return nil, t.cfg.getError
   142  	}
   143  	return t.Snapshot.BatchGet(ctx, keys)
   144  }