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

     1  package rawkv
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"math/rand"
     8  	"strconv"
     9  	"time"
    10  
    11  	"github.com/anishathalye/porcupine"
    12  	"github.com/pingcap/chaos/pkg/core"
    13  	"github.com/pingcap/chaos/pkg/model"
    14  	"github.com/pingcap/tidb/config"
    15  	"github.com/pingcap/tidb/store/tikv"
    16  )
    17  
    18  var (
    19  	register = []byte("acc")
    20  )
    21  
    22  type registerClient struct {
    23  	db *tikv.RawKVClient
    24  	r  *rand.Rand
    25  }
    26  
    27  func (c *registerClient) SetUp(ctx context.Context, nodes []string, node string) error {
    28  	c.r = rand.New(rand.NewSource(time.Now().UnixNano()))
    29  	db, err := tikv.NewRawKVClient([]string{fmt.Sprintf("%s:2379", node)}, config.Security{})
    30  	if err != nil {
    31  		return err
    32  	}
    33  
    34  	c.db = db
    35  
    36  	// Do SetUp in the first node
    37  	if node != nodes[0] {
    38  		return nil
    39  	}
    40  
    41  	log.Printf("begin to initial register on node %s", node)
    42  
    43  	db.Put(register, []byte("0"))
    44  
    45  	return nil
    46  }
    47  
    48  func (c *registerClient) TearDown(ctx context.Context, nodes []string, node string) error {
    49  	return c.db.Close()
    50  }
    51  
    52  func (c *registerClient) invokeRead(ctx context.Context, r model.RegisterRequest) model.RegisterResponse {
    53  	val, err := c.db.Get(register)
    54  	if err != nil {
    55  		return model.RegisterResponse{Unknown: true}
    56  	}
    57  	v, err := strconv.ParseInt(string(val), 10, 64)
    58  	if err != nil {
    59  		panic(fmt.Sprintf("invalid value: %s", val))
    60  	}
    61  	return model.RegisterResponse{Value: int(v)}
    62  }
    63  
    64  func (c *registerClient) Invoke(ctx context.Context, node string, r interface{}) interface{} {
    65  	arg := r.(model.RegisterRequest)
    66  	if arg.Op == model.RegisterRead {
    67  		return c.invokeRead(ctx, arg)
    68  	}
    69  
    70  	val := fmt.Sprintf("%d", arg.Value)
    71  	err := c.db.Put(register, []byte(val))
    72  	if err != nil {
    73  		return model.RegisterResponse{Unknown: true}
    74  	}
    75  	return model.RegisterResponse{}
    76  }
    77  
    78  func (c *registerClient) NextRequest() interface{} {
    79  	r := model.RegisterRequest{
    80  		Op: c.r.Intn(2) == 1,
    81  	}
    82  	if r.Op == model.RegisterRead {
    83  		return r
    84  	}
    85  
    86  	r.Value = int(c.r.Int63())
    87  	return r
    88  }
    89  
    90  // DumpState the database state(also the model's state)
    91  func (c *registerClient) DumpState(ctx context.Context) (interface{}, error) {
    92  	val, err := c.db.Get(register)
    93  	if err != nil {
    94  		return nil, err
    95  	}
    96  	v, err := strconv.ParseInt(string(val), 10, 64)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  	return v, nil
   101  }
   102  
   103  func newRegisterEvent(v interface{}, id uint) porcupine.Event {
   104  	if _, ok := v.(model.RegisterRequest); ok {
   105  		return porcupine.Event{Kind: porcupine.CallEvent, Value: v, Id: id}
   106  	}
   107  
   108  	return porcupine.Event{Kind: porcupine.ReturnEvent, Value: v, Id: id}
   109  }
   110  
   111  // RegisterClientCreator creates a register test client for rawkv.
   112  type RegisterClientCreator struct {
   113  }
   114  
   115  // Create creates a client.
   116  func (RegisterClientCreator) Create(node string) core.Client {
   117  	return &registerClient{}
   118  }