github.com/pingcap/chaos@v0.0.0-20190710112158-c86faf4b3719/db/txnkv/register.go (about)

     1  package txnkv
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"math/rand"
     8  	"strconv"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/anishathalye/porcupine"
    13  	"github.com/pingcap/chaos/pkg/core"
    14  	"github.com/pingcap/chaos/pkg/model"
    15  	"github.com/pingcap/tidb/kv"
    16  	"github.com/pingcap/tidb/store/tikv"
    17  )
    18  
    19  var (
    20  	register = []byte("acc")
    21  
    22  	closeOnce = sync.Once{}
    23  )
    24  
    25  type registerClient struct {
    26  	db kv.Storage
    27  	r  *rand.Rand
    28  }
    29  
    30  func (c *registerClient) SetUp(ctx context.Context, nodes []string, node string) error {
    31  	c.r = rand.New(rand.NewSource(time.Now().UnixNano()))
    32  	driver := tikv.Driver{}
    33  	db, err := driver.Open(fmt.Sprintf("tikv://%s:2379?disableGC=true", node))
    34  	if err != nil {
    35  		return err
    36  	}
    37  
    38  	c.db = db
    39  
    40  	// Do SetUp in the first node
    41  	if node != nodes[0] {
    42  		return nil
    43  	}
    44  
    45  	log.Printf("begin to initial register on node %s", node)
    46  
    47  	tx, err := db.Begin()
    48  	if err != nil {
    49  		return err
    50  	}
    51  
    52  	defer tx.Rollback()
    53  
    54  	if err = tx.Set(register, []byte("0")); err != nil {
    55  		return err
    56  	}
    57  
    58  	return tx.Commit(ctx)
    59  }
    60  
    61  func (c *registerClient) TearDown(ctx context.Context, nodes []string, node string) error {
    62  	var err error
    63  	closeOnce.Do(func() {
    64  		// It's a workaround for `panic: close of closed channel`.
    65  		// `tikv.Driver.Open` will open the same instance if cluster id is
    66  		// the same.
    67  		//
    68  		// See more: https://github.com/pingcap/tidb/blob/
    69  		//           63c4562c27ad43165e6a0d5f890f33f3b1002b3f/store/tikv/kv.go#L95
    70  		err = c.db.Close()
    71  	})
    72  	return err
    73  }
    74  
    75  func (c *registerClient) invokeRead(ctx context.Context, r model.RegisterRequest) model.RegisterResponse {
    76  	tx, err := c.db.Begin()
    77  	if err != nil {
    78  		return model.RegisterResponse{Unknown: true}
    79  	}
    80  	defer tx.Rollback()
    81  
    82  	val, err := tx.Get(register)
    83  	if err != nil {
    84  		return model.RegisterResponse{Unknown: true}
    85  	}
    86  
    87  	if err = tx.Commit(ctx); err != nil {
    88  		return model.RegisterResponse{Unknown: true}
    89  	}
    90  
    91  	v, err := strconv.ParseInt(string(val), 10, 64)
    92  	if err != nil {
    93  		return model.RegisterResponse{Unknown: true}
    94  	}
    95  	return model.RegisterResponse{Value: int(v)}
    96  }
    97  
    98  func (c *registerClient) invokeWrite(ctx context.Context, r model.RegisterRequest) model.RegisterResponse {
    99  	tx, err := c.db.Begin()
   100  	if err != nil {
   101  		return model.RegisterResponse{Unknown: true}
   102  	}
   103  	defer tx.Rollback()
   104  
   105  	val := fmt.Sprintf("%d", r.Value)
   106  	if err = tx.Set(register, []byte(val)); err != nil {
   107  		return model.RegisterResponse{Unknown: true}
   108  	}
   109  
   110  	if err = tx.Commit(ctx); err != nil {
   111  		return model.RegisterResponse{Unknown: true}
   112  	}
   113  	return model.RegisterResponse{}
   114  }
   115  
   116  func (c *registerClient) Invoke(ctx context.Context, node string, r interface{}) interface{} {
   117  	arg := r.(model.RegisterRequest)
   118  	if arg.Op == model.RegisterRead {
   119  		return c.invokeRead(ctx, arg)
   120  	}
   121  	return c.invokeWrite(ctx, arg)
   122  }
   123  
   124  func (c *registerClient) NextRequest() interface{} {
   125  	r := model.RegisterRequest{
   126  		Op: c.r.Intn(2) == 1,
   127  	}
   128  	if r.Op == model.RegisterRead {
   129  		return r
   130  	}
   131  
   132  	r.Value = int(c.r.Int63())
   133  	return r
   134  }
   135  
   136  // DumpState the database state(also the model's state)
   137  func (c *registerClient) DumpState(ctx context.Context) (interface{}, error) {
   138  	tx, err := c.db.Begin()
   139  	if err != nil {
   140  		return nil, err
   141  	}
   142  	defer tx.Rollback()
   143  
   144  	val, err := tx.Get(register)
   145  	if err != nil {
   146  		return nil, err
   147  	}
   148  
   149  	if err = tx.Commit(ctx); err != nil {
   150  		return nil, err
   151  	}
   152  
   153  	v, err := strconv.ParseInt(string(val), 10, 64)
   154  	if err != nil {
   155  		return nil, err
   156  	}
   157  	return v, nil
   158  }
   159  
   160  func newRegisterEvent(v interface{}, id uint) porcupine.Event {
   161  	if _, ok := v.(model.RegisterRequest); ok {
   162  		return porcupine.Event{Kind: porcupine.CallEvent, Value: v, Id: id}
   163  	}
   164  
   165  	return porcupine.Event{Kind: porcupine.ReturnEvent, Value: v, Id: id}
   166  }
   167  
   168  // RegisterClientCreator creates a register test client for txnkv.
   169  type RegisterClientCreator struct {
   170  }
   171  
   172  // Create creates a client.
   173  func (RegisterClientCreator) Create(node string) core.Client {
   174  	return &registerClient{}
   175  }