github.com/matrixorigin/matrixone@v0.7.0/pkg/txn/service/service_cn_handler_test.go (about)

     1  // Copyright 2021 - 2022 Matrix Origin
     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 service
    16  
    17  import (
    18  	"context"
    19  	"math"
    20  	"sync/atomic"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    25  	"github.com/matrixorigin/matrixone/pkg/pb/timestamp"
    26  	"github.com/matrixorigin/matrixone/pkg/pb/txn"
    27  	"github.com/matrixorigin/matrixone/pkg/txn/rpc"
    28  	"github.com/matrixorigin/matrixone/pkg/txn/storage/mem"
    29  	"github.com/stretchr/testify/assert"
    30  	"github.com/stretchr/testify/require"
    31  )
    32  
    33  func TestReadBasic(t *testing.T) {
    34  	sender := NewTestSender()
    35  	defer func() {
    36  		assert.NoError(t, sender.Close())
    37  	}()
    38  
    39  	s := NewTestTxnService(t, 1, sender, NewTestClock(0))
    40  	assert.NoError(t, s.Start())
    41  	defer func() {
    42  		assert.NoError(t, s.Close(false))
    43  	}()
    44  
    45  	sender.AddTxnService(s)
    46  
    47  	rTxn := NewTestTxn(1, 1)
    48  	resp := readTestData(t, sender, 1, rTxn, 1)
    49  	checkReadResponses(t, resp, "")
    50  }
    51  
    52  func TestReadWithDNShardNotMatch(t *testing.T) {
    53  	sender := NewTestSender()
    54  	defer func() {
    55  		assert.NoError(t, sender.Close())
    56  	}()
    57  	sender.setFilter(func(req *txn.TxnRequest) bool {
    58  		req.CNRequest.Target.ReplicaID = 0
    59  		return true
    60  	})
    61  
    62  	s := NewTestTxnService(t, 1, sender, NewTestClock(0))
    63  	assert.NoError(t, s.Start())
    64  	defer func() {
    65  		assert.NoError(t, s.Close(false))
    66  	}()
    67  
    68  	sender.AddTxnService(s)
    69  
    70  	rTxn := NewTestTxn(1, 1)
    71  	resp := readTestData(t, sender, 1, rTxn, 1)
    72  	checkResponses(t, resp,
    73  		txn.WrapError(moerr.NewDNShardNotFound(context.TODO(), "", 1), 0))
    74  	// newTxnError(moerr.ErrDNShardNotFound, "txn not active"))
    75  }
    76  
    77  func TestReadWithSelfWrite(t *testing.T) {
    78  	sender := NewTestSender()
    79  	defer func() {
    80  		assert.NoError(t, sender.Close())
    81  	}()
    82  
    83  	s := NewTestTxnService(t, 1, sender, NewTestClock(0))
    84  	assert.NoError(t, s.Start())
    85  	defer func() {
    86  		assert.NoError(t, s.Close(false))
    87  	}()
    88  	sender.AddTxnService(s)
    89  
    90  	rwTxn := NewTestTxn(1, 1)
    91  	checkResponses(t, writeTestData(t, sender, 1, rwTxn, 1))
    92  	checkReadResponses(t, readTestData(t, sender, 1, rwTxn, 1), "1-1-1")
    93  	checkResponses(t, writeTestData(t, sender, 1, rwTxn, 2))
    94  	checkReadResponses(t, readTestData(t, sender, 1, rwTxn, 2), "2-1-1")
    95  }
    96  
    97  func TestReadBlockWithClock(t *testing.T) {
    98  	sender := NewTestSender()
    99  	defer func() {
   100  		assert.NoError(t, sender.Close())
   101  	}()
   102  
   103  	ts := int64(0)
   104  	s := NewTestTxnService(t, 1, sender, NewTestSpecClock(func() int64 {
   105  		return atomic.AddInt64(&ts, 1)
   106  	}))
   107  	assert.NoError(t, s.Start())
   108  	defer func() {
   109  		assert.NoError(t, s.Close(false))
   110  	}()
   111  	sender.AddTxnService(s)
   112  
   113  	c := make(chan struct{})
   114  	go func() {
   115  		rwTxn := NewTestTxn(1, 3)
   116  		checkReadResponses(t, readTestData(t, sender, 1, rwTxn, 1), "")
   117  		c <- struct{}{}
   118  	}()
   119  	<-c
   120  	assert.Equal(t, int64(3), ts)
   121  }
   122  
   123  func TestReadCannotBlockByUncomitted(t *testing.T) {
   124  	sender := NewTestSender()
   125  	defer func() {
   126  		assert.NoError(t, sender.Close())
   127  	}()
   128  
   129  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   130  	assert.NoError(t, s.Start())
   131  	defer func() {
   132  		assert.NoError(t, s.Close(false))
   133  	}()
   134  	sender.AddTxnService(s)
   135  
   136  	wTxn := NewTestTxn(1, 1)
   137  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   138  
   139  	rTxn := NewTestTxn(2, 1)
   140  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   141  }
   142  
   143  func TestReadCannotBlockByPreparedIfSnapshotTSIsLEPreparedTS(t *testing.T) {
   144  	sender := NewTestSender()
   145  	defer func() {
   146  		assert.NoError(t, sender.Close())
   147  	}()
   148  
   149  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   150  	assert.NoError(t, s.Start())
   151  	defer func() {
   152  		assert.NoError(t, s.Close(false))
   153  	}()
   154  	sender.AddTxnService(s)
   155  
   156  	wTxn := NewTestTxn(1, 1)
   157  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   158  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   159  
   160  	rTxn := NewTestTxn(2, 1)
   161  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   162  
   163  	rTxn = NewTestTxn(2, 2)
   164  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   165  }
   166  
   167  func TestReadWillBlockByPreparedIfSnapshotTSIsGTPreparedTS(t *testing.T) {
   168  	sender := NewTestSender()
   169  	defer func() {
   170  		assert.NoError(t, sender.Close())
   171  	}()
   172  
   173  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   174  	assert.NoError(t, s.Start())
   175  	defer func() {
   176  		assert.NoError(t, s.Close(false))
   177  	}()
   178  	sender.AddTxnService(s)
   179  
   180  	wTxn := NewTestTxn(1, 1)
   181  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   182  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   183  
   184  	c := make(chan struct{})
   185  	go func() {
   186  		rTxn := NewTestTxn(2, 3)
   187  		readTestData(t, sender, 1, rTxn, 1)
   188  		close(c)
   189  	}()
   190  	select {
   191  	case <-c:
   192  		assert.Fail(t, "cannot read")
   193  	case <-time.After(time.Second):
   194  
   195  	}
   196  }
   197  
   198  func TestReadAfterBlockTxnCommitted(t *testing.T) {
   199  	sender := NewTestSender()
   200  	defer func() {
   201  		assert.NoError(t, sender.Close())
   202  	}()
   203  
   204  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   205  	assert.NoError(t, s.Start())
   206  	defer func() {
   207  		assert.NoError(t, s.Close(false))
   208  	}()
   209  	sender.AddTxnService(s)
   210  
   211  	wTxn := NewTestTxn(1, 1, 1)
   212  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   213  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   214  
   215  	c := make(chan struct{})
   216  	go func() {
   217  		rTxn := NewTestTxn(2, 3)
   218  		checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "1-1-1")
   219  		close(c)
   220  	}()
   221  	go func() {
   222  		time.Sleep(time.Second)
   223  		wTxn.CommitTS = NewTestTimestamp(2) // commit at 2
   224  		checkResponses(t, commitShardWriteData(t, sender, wTxn))
   225  	}()
   226  
   227  	select {
   228  	case <-c:
   229  	case <-time.After(time.Minute):
   230  		assert.Fail(t, "cannot read")
   231  	}
   232  }
   233  
   234  func TestReadAfterBlockTxnCommittedAndCannotReadCommittedValue(t *testing.T) {
   235  	sender := NewTestSender()
   236  	defer func() {
   237  		assert.NoError(t, sender.Close())
   238  	}()
   239  
   240  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   241  	assert.NoError(t, s.Start())
   242  	defer func() {
   243  		assert.NoError(t, s.Close(false))
   244  	}()
   245  	sender.AddTxnService(s)
   246  
   247  	wTxn := NewTestTxn(1, 1, 1)
   248  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   249  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   250  
   251  	c := make(chan struct{})
   252  	go func() {
   253  		rTxn := NewTestTxn(2, 3)
   254  		checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   255  		close(c)
   256  	}()
   257  	go func() {
   258  		time.Sleep(time.Second)
   259  		wTxn.CommitTS = NewTestTimestamp(3) // commit at 3
   260  		checkResponses(t, commitShardWriteData(t, sender, wTxn))
   261  	}()
   262  
   263  	select {
   264  	case <-c:
   265  	case <-time.After(time.Minute):
   266  		assert.Fail(t, "cannot read")
   267  	}
   268  }
   269  
   270  func TestReadAfterBlockTxnAborted(t *testing.T) {
   271  	sender := NewTestSender()
   272  	defer func() {
   273  		assert.NoError(t, sender.Close())
   274  	}()
   275  
   276  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   277  	assert.NoError(t, s.Start())
   278  	defer func() {
   279  		assert.NoError(t, s.Close(false))
   280  	}()
   281  	sender.AddTxnService(s)
   282  
   283  	wTxn := NewTestTxn(1, 1, 1)
   284  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   285  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   286  
   287  	c := make(chan struct{})
   288  	go func() {
   289  		rTxn := NewTestTxn(2, 3)
   290  		checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   291  		close(c)
   292  	}()
   293  	go func() {
   294  		time.Sleep(time.Second)
   295  		checkResponses(t, rollbackShardWriteData(t, sender, wTxn))
   296  	}()
   297  
   298  	select {
   299  	case <-c:
   300  	case <-time.After(time.Minute):
   301  		assert.Fail(t, "cannot read")
   302  	}
   303  }
   304  
   305  func TestReadCannotBlockByCommittingIfSnapshotTSIsLECommitTS(t *testing.T) {
   306  	sender := NewTestSender()
   307  	defer func() {
   308  		assert.NoError(t, sender.Close())
   309  	}()
   310  
   311  	s := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   312  	assert.NoError(t, s.Start())
   313  	defer func() {
   314  		assert.NoError(t, s.Close(false))
   315  	}()
   316  	sender.AddTxnService(s)
   317  
   318  	wTxn := NewTestTxn(1, 1)
   319  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   320  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   321  
   322  	wTxn.CommitTS = NewTestTimestamp(2)
   323  	assert.NoError(t, s.storage.(*mem.KVTxnStorage).Committing(context.TODO(), wTxn))
   324  
   325  	rTxn := NewTestTxn(2, 1)
   326  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   327  
   328  	rTxn = NewTestTxn(2, 2)
   329  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   330  }
   331  
   332  func TestReadWillBlockByCommittingIfSnapshotTSIsGTCommitTS(t *testing.T) {
   333  	sender := NewTestSender()
   334  	defer func() {
   335  		assert.NoError(t, sender.Close())
   336  	}()
   337  
   338  	s := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   339  	assert.NoError(t, s.Start())
   340  	defer func() {
   341  		assert.NoError(t, s.Close(false))
   342  	}()
   343  	sender.AddTxnService(s)
   344  
   345  	wTxn := NewTestTxn(1, 1)
   346  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   347  	prepareTestTxn(t, sender, wTxn, 1) // prepare at 2
   348  
   349  	wTxn.CommitTS = NewTestTimestamp(2)
   350  	assert.NoError(t, s.storage.(*mem.KVTxnStorage).Committing(context.TODO(), wTxn))
   351  
   352  	c := make(chan struct{})
   353  	go func() {
   354  		rTxn := NewTestTxn(2, 3)
   355  		readTestData(t, sender, 1, rTxn, 1)
   356  		close(c)
   357  	}()
   358  	select {
   359  	case <-c:
   360  		assert.Fail(t, "cannot read")
   361  	case <-time.After(time.Second):
   362  
   363  	}
   364  }
   365  
   366  func TestReadCommitted(t *testing.T) {
   367  	sender := NewTestSender()
   368  	defer func() {
   369  		assert.NoError(t, sender.Close())
   370  	}()
   371  
   372  	s := NewTestTxnService(t, 1, sender, NewTestClock(1))
   373  	assert.NoError(t, s.Start())
   374  	defer func() {
   375  		assert.NoError(t, s.Close(false))
   376  	}()
   377  	sender.AddTxnService(s)
   378  
   379  	wTxn1 := NewTestTxn(1, 1, 1) // commit at 2
   380  	checkResponses(t, writeTestData(t, sender, 1, wTxn1, 1))
   381  	checkResponses(t, commitWriteData(t, sender, wTxn1))
   382  
   383  	wTxn2 := NewTestTxn(2, 1, 1) // commit at 3
   384  	checkResponses(t, writeTestData(t, sender, 1, wTxn2, 2))
   385  	checkResponses(t, commitWriteData(t, sender, wTxn2))
   386  
   387  	rTxn := NewTestTxn(3, 2)
   388  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), "")
   389  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 2), "")
   390  
   391  	rTxn = NewTestTxn(3, 3)
   392  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), string(GetTestValue(1, wTxn1)))
   393  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 2), "")
   394  
   395  	rTxn = NewTestTxn(3, 4)
   396  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 1), string(GetTestValue(1, wTxn1)))
   397  	checkReadResponses(t, readTestData(t, sender, 1, rTxn, 2), string(GetTestValue(2, wTxn2)))
   398  }
   399  
   400  func TestWriteBasic(t *testing.T) {
   401  	sender := NewTestSender()
   402  	defer func() {
   403  		assert.NoError(t, sender.Close())
   404  	}()
   405  
   406  	s := NewTestTxnService(t, 1, sender, NewTestClock(0)).(*service)
   407  	assert.NoError(t, s.Start())
   408  	defer func() {
   409  		assert.NoError(t, s.Close(false))
   410  	}()
   411  
   412  	sender.AddTxnService(s)
   413  
   414  	wTxn := NewTestTxn(1, 1)
   415  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   416  
   417  	kv := s.storage.(*mem.KVTxnStorage).GetUncommittedKV()
   418  	v, ok := kv.Get(GetTestKey(1))
   419  	assert.True(t, ok)
   420  	assert.Equal(t, GetTestValue(1, wTxn), v)
   421  }
   422  
   423  func TestWriteWithDNShardNotMatch(t *testing.T) {
   424  	sender := NewTestSender()
   425  	defer func() {
   426  		assert.NoError(t, sender.Close())
   427  	}()
   428  	sender.setFilter(func(req *txn.TxnRequest) bool {
   429  		req.CNRequest.Target.ReplicaID = 0
   430  		return true
   431  	})
   432  
   433  	s := NewTestTxnService(t, 1, sender, NewTestClock(0)).(*service)
   434  	assert.NoError(t, s.Start())
   435  	defer func() {
   436  		assert.NoError(t, s.Close(false))
   437  	}()
   438  
   439  	sender.AddTxnService(s)
   440  
   441  	wTxn := NewTestTxn(1, 1)
   442  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1),
   443  		txn.WrapError(moerr.NewDNShardNotFound(context.TODO(), "", 1), 0))
   444  }
   445  
   446  func TestWriteWithWWConflict(t *testing.T) {
   447  	sender := NewTestSender()
   448  	defer func() {
   449  		assert.NoError(t, sender.Close())
   450  	}()
   451  
   452  	s := NewTestTxnService(t, 1, sender, NewTestClock(0))
   453  	assert.NoError(t, s.Start())
   454  	defer func() {
   455  		assert.NoError(t, s.Close(false))
   456  	}()
   457  
   458  	sender.AddTxnService(s)
   459  
   460  	wTxn := NewTestTxn(1, 1)
   461  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   462  
   463  	wTxn2 := NewTestTxn(2, 1)
   464  	checkResponses(t, writeTestData(t, sender, 1, wTxn2, 1),
   465  		txn.WrapError(moerr.NewTAEWrite(context.TODO()), 0))
   466  	// newTxnError(moerr.ErrTAEWrite, "write conlict"))
   467  }
   468  
   469  func TestCommitWithSingleDNShard(t *testing.T) {
   470  	sender := NewTestSender()
   471  	defer func() {
   472  		assert.NoError(t, sender.Close())
   473  	}()
   474  
   475  	s := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   476  	assert.NoError(t, s.Start())
   477  	defer func() {
   478  		assert.NoError(t, s.Close(false))
   479  	}()
   480  
   481  	sender.AddTxnService(s)
   482  
   483  	n := byte(10)
   484  	wTxn := NewTestTxn(1, 1, 1)
   485  	for i := byte(0); i < n; i++ {
   486  		checkResponses(t, writeTestData(t, sender, 1, wTxn, i))
   487  	}
   488  	checkResponses(t, commitWriteData(t, sender, wTxn))
   489  
   490  	for i := byte(0); i < n; i++ {
   491  		var values [][]byte
   492  		var timestamps []timestamp.Timestamp
   493  		kv := s.storage.(*mem.KVTxnStorage).GetCommittedKV()
   494  
   495  		kv.AscendRange(GetTestKey(i), NewTestTimestamp(0), NewTestTimestamp(math.MaxInt64), func(value []byte, ts timestamp.Timestamp) {
   496  			values = append(values, value)
   497  			timestamps = append(timestamps, ts)
   498  		})
   499  		assert.Equal(t, [][]byte{GetTestValue(i, wTxn)}, values)
   500  		assert.Equal(t, []timestamp.Timestamp{NewTestTimestamp(2)}, timestamps)
   501  	}
   502  }
   503  
   504  func TestCommitWithDNShardNotMatch(t *testing.T) {
   505  	sender := NewTestSender()
   506  	defer func() {
   507  		assert.NoError(t, sender.Close())
   508  	}()
   509  	sender.setFilter(func(req *txn.TxnRequest) bool {
   510  		if req.CommitRequest != nil {
   511  			req.Txn.DNShards[0].ReplicaID = 0
   512  		}
   513  		return true
   514  	})
   515  
   516  	s := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   517  	assert.NoError(t, s.Start())
   518  	defer func() {
   519  		assert.NoError(t, s.Close(false))
   520  	}()
   521  
   522  	sender.AddTxnService(s)
   523  
   524  	n := byte(10)
   525  	wTxn := NewTestTxn(1, 1, 1)
   526  	for i := byte(0); i < n; i++ {
   527  		checkResponses(t, writeTestData(t, sender, 1, wTxn, i))
   528  	}
   529  	checkResponses(t, commitWriteData(t, sender, wTxn),
   530  		txn.WrapError(moerr.NewDNShardNotFound(context.TODO(), "", 1), 0))
   531  }
   532  
   533  func TestCommitWithMultiDNShards(t *testing.T) {
   534  	sender := NewTestSender()
   535  	defer func() {
   536  		assert.NoError(t, sender.Close())
   537  	}()
   538  
   539  	s1 := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   540  	assert.NoError(t, s1.Start())
   541  	defer func() {
   542  		assert.NoError(t, s1.Close(false))
   543  	}()
   544  	s2 := NewTestTxnService(t, 2, sender, NewTestClock(1)).(*service)
   545  	assert.NoError(t, s2.Start())
   546  	defer func() {
   547  		assert.NoError(t, s2.Close(false))
   548  	}()
   549  
   550  	sender.AddTxnService(s1)
   551  	sender.AddTxnService(s2)
   552  
   553  	wTxn := NewTestTxn(1, 1, 1, 2)
   554  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   555  	checkResponses(t, writeTestData(t, sender, 2, wTxn, 2))
   556  
   557  	w1 := addTestWaiter(t, s1, wTxn, txn.TxnStatus_Committed)
   558  	defer w1.close()
   559  	w2 := addTestWaiter(t, s2, wTxn, txn.TxnStatus_Committed)
   560  	defer w2.close()
   561  
   562  	checkResponses(t, commitWriteData(t, sender, wTxn))
   563  
   564  	checkWaiter(t, w1, txn.TxnStatus_Committed)
   565  	checkWaiter(t, w2, txn.TxnStatus_Committed)
   566  
   567  	checkData(t, wTxn, s1, 2, 1, true)
   568  	checkData(t, wTxn, s2, 2, 2, true)
   569  }
   570  
   571  func TestCommitWithRollbackIfAnyPrepareFailed(t *testing.T) {
   572  	sender := NewTestSender()
   573  	defer func() {
   574  		assert.NoError(t, sender.Close())
   575  	}()
   576  
   577  	s1 := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   578  	assert.NoError(t, s1.Start())
   579  	defer func() {
   580  		assert.NoError(t, s1.Close(false))
   581  	}()
   582  	s2 := NewTestTxnService(t, 2, sender, NewTestClock(1)).(*service)
   583  	assert.NoError(t, s2.Start())
   584  	defer func() {
   585  		assert.NoError(t, s2.Close(false))
   586  	}()
   587  
   588  	sender.AddTxnService(s1)
   589  	sender.AddTxnService(s2)
   590  
   591  	wTxn1 := NewTestTxn(1, 1, 1)
   592  	writeTestData(t, sender, 1, wTxn1, 1)
   593  	checkResponses(t, commitWriteData(t, sender, wTxn1)) // commit at 2
   594  
   595  	wTxn := NewTestTxn(1, 1, 1, 2)
   596  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   597  	checkResponses(t, writeTestData(t, sender, 2, wTxn, 2))
   598  
   599  	w1 := addTestWaiter(t, s1, wTxn, txn.TxnStatus_Aborted)
   600  	defer w1.close()
   601  	w2 := addTestWaiter(t, s2, wTxn, txn.TxnStatus_Aborted)
   602  	defer w2.close()
   603  
   604  	checkResponses(t, commitWriteData(t, sender, wTxn),
   605  		txn.WrapError(moerr.NewTAEPrepare(context.TODO(), "cannot prepare"), 0))
   606  
   607  	checkWaiter(t, w1, txn.TxnStatus_Aborted)
   608  	checkWaiter(t, w2, txn.TxnStatus_Aborted)
   609  
   610  	checkData(t, wTxn, s1, 2, 2, false)
   611  	checkData(t, wTxn, s2, 2, 2, false)
   612  }
   613  
   614  func TestRollback(t *testing.T) {
   615  	sender := NewTestSender()
   616  	defer func() {
   617  		assert.NoError(t, sender.Close())
   618  	}()
   619  
   620  	s1 := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   621  	assert.NoError(t, s1.Start())
   622  	defer func() {
   623  		assert.NoError(t, s1.Close(false))
   624  	}()
   625  	s2 := NewTestTxnService(t, 2, sender, NewTestClock(1)).(*service)
   626  	assert.NoError(t, s2.Start())
   627  	defer func() {
   628  		assert.NoError(t, s2.Close(false))
   629  	}()
   630  
   631  	sender.AddTxnService(s1)
   632  	sender.AddTxnService(s2)
   633  
   634  	wTxn := NewTestTxn(1, 1, 1)
   635  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   636  	wTxn.DNShards = append(wTxn.DNShards, NewTestDNShard(2))
   637  	checkResponses(t, writeTestData(t, sender, 2, wTxn, 2))
   638  
   639  	w1 := addTestWaiter(t, s1, wTxn, txn.TxnStatus_Aborted)
   640  	defer w1.close()
   641  	w2 := addTestWaiter(t, s2, wTxn, txn.TxnStatus_Aborted)
   642  	defer w2.close()
   643  
   644  	responses := rollbackWriteData(t, sender, wTxn)
   645  	checkResponses(t, responses)
   646  	for _, resp := range responses {
   647  		assert.Equal(t, txn.TxnStatus_Aborted, resp.Txn.Status)
   648  	}
   649  
   650  	checkWaiter(t, w1, txn.TxnStatus_Aborted)
   651  	checkWaiter(t, w2, txn.TxnStatus_Aborted)
   652  
   653  	checkData(t, wTxn, s1, 2, 0, false)
   654  	checkData(t, wTxn, s2, 2, 0, false)
   655  }
   656  
   657  func TestRollbackWithDNShardNotFound(t *testing.T) {
   658  	sender := NewTestSender()
   659  	defer func() {
   660  		assert.NoError(t, sender.Close())
   661  	}()
   662  
   663  	sender.setFilter(func(req *txn.TxnRequest) bool {
   664  		if req.RollbackRequest != nil {
   665  			req.Txn.DNShards[0].ReplicaID = 0
   666  		}
   667  		return true
   668  	})
   669  
   670  	s1 := NewTestTxnService(t, 1, sender, NewTestClock(1)).(*service)
   671  	assert.NoError(t, s1.Start())
   672  	defer func() {
   673  		assert.NoError(t, s1.Close(false))
   674  	}()
   675  	s2 := NewTestTxnService(t, 2, sender, NewTestClock(1)).(*service)
   676  	assert.NoError(t, s2.Start())
   677  	defer func() {
   678  		assert.NoError(t, s2.Close(false))
   679  	}()
   680  
   681  	sender.AddTxnService(s1)
   682  	sender.AddTxnService(s2)
   683  
   684  	wTxn := NewTestTxn(1, 1, 1)
   685  	checkResponses(t, writeTestData(t, sender, 1, wTxn, 1))
   686  	wTxn.DNShards = append(wTxn.DNShards, NewTestDNShard(2))
   687  	checkResponses(t, writeTestData(t, sender, 2, wTxn, 2))
   688  
   689  	checkResponses(t, rollbackWriteData(t, sender, wTxn),
   690  		txn.WrapError(moerr.NewDNShardNotFound(context.TODO(), "", 1), 0))
   691  }
   692  
   693  func writeTestData(t *testing.T, sender rpc.TxnSender, toShard uint64, wTxn txn.TxnMeta, keys ...byte) []txn.TxnResponse {
   694  	requests := make([]txn.TxnRequest, 0, len(keys))
   695  	for _, k := range keys {
   696  		requests = append(requests, NewTestWriteRequest(k, wTxn, toShard))
   697  	}
   698  	result, err := sender.Send(context.Background(), requests)
   699  	assert.NoError(t, err)
   700  	responses := result.Responses
   701  	assert.Equal(t, len(keys), len(responses))
   702  	return responses
   703  }
   704  
   705  func TestDebug(t *testing.T) {
   706  	sender := NewTestSender()
   707  	defer func() {
   708  		assert.NoError(t, sender.Close())
   709  	}()
   710  
   711  	s := NewTestTxnService(t, 1, sender, NewTestClock(0))
   712  	assert.NoError(t, s.Start())
   713  	defer func() {
   714  		assert.NoError(t, s.Close(false))
   715  	}()
   716  
   717  	sender.AddTxnService(s)
   718  
   719  	data := []byte("OK")
   720  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
   721  	defer cancel()
   722  
   723  	result, err := sender.Send(ctx, []txn.TxnRequest{
   724  		{
   725  			Method: txn.TxnMethod_DEBUG,
   726  			CNRequest: &txn.CNOpRequest{
   727  				Payload: data,
   728  				Target:  NewTestDNShard(1),
   729  			},
   730  		},
   731  	})
   732  	require.NoError(t, err)
   733  	assert.Equal(t, data, result.Responses[0].CNOpResponse.Payload)
   734  }
   735  
   736  func commitShardWriteData(t *testing.T, sender rpc.TxnSender, wTxn txn.TxnMeta) []txn.TxnResponse {
   737  	result, err := sender.Send(context.Background(), []txn.TxnRequest{NewTestCommitShardRequest(wTxn)})
   738  	assert.NoError(t, err)
   739  	responses := result.Responses
   740  	return responses
   741  }
   742  
   743  func rollbackShardWriteData(t *testing.T, sender rpc.TxnSender, wTxn txn.TxnMeta) []txn.TxnResponse {
   744  	result, err := sender.Send(context.Background(), []txn.TxnRequest{NewTestRollbackShardRequest(wTxn)})
   745  	assert.NoError(t, err)
   746  	responses := result.Responses
   747  	return responses
   748  }
   749  
   750  func commitWriteData(t *testing.T, sender rpc.TxnSender, wTxn txn.TxnMeta) []txn.TxnResponse {
   751  	result, err := sender.Send(context.Background(), []txn.TxnRequest{NewTestCommitRequest(wTxn)})
   752  	assert.NoError(t, err)
   753  	responses := result.Responses
   754  	return responses
   755  }
   756  
   757  func rollbackWriteData(t *testing.T, sender rpc.TxnSender, wTxn txn.TxnMeta) []txn.TxnResponse {
   758  	result, err := sender.Send(context.Background(), []txn.TxnRequest{NewTestRollbackRequest(wTxn)})
   759  	assert.NoError(t, err)
   760  	responses := result.Responses
   761  	return responses
   762  }
   763  
   764  func readTestData(t *testing.T, sender rpc.TxnSender, toShard uint64, rTxn txn.TxnMeta, keys ...byte) []txn.TxnResponse {
   765  	requests := make([]txn.TxnRequest, 0, len(keys))
   766  	for _, k := range keys {
   767  		requests = append(requests, NewTestReadRequest(k, rTxn, toShard))
   768  	}
   769  	result, err := sender.Send(context.Background(), requests)
   770  	assert.NoError(t, err)
   771  	responses := result.Responses
   772  	assert.Equal(t, len(keys), len(responses))
   773  	return responses
   774  }
   775  
   776  func checkReadResponses(t *testing.T, response []txn.TxnResponse, expectValues ...string) {
   777  	for idx, resp := range response {
   778  		values := mem.MustParseGetPayload(resp.CNOpResponse.Payload)
   779  		assert.Equal(t, expectValues[idx], string(values[0]))
   780  		assert.NotNil(t, resp.Txn)
   781  	}
   782  }
   783  
   784  func checkResponses(t *testing.T, response []txn.TxnResponse, expectErrors ...*txn.TxnError) {
   785  	if len(expectErrors) == 0 {
   786  		expectErrors = make([]*txn.TxnError, len(response))
   787  	}
   788  	for idx, resp := range response {
   789  		if resp.TxnError == nil {
   790  			assert.Equal(t, expectErrors[idx], resp.TxnError)
   791  		} else {
   792  			assert.Equal(t, expectErrors[idx].TxnErrCode, resp.TxnError.TxnErrCode)
   793  		}
   794  	}
   795  }
   796  
   797  func checkData(t *testing.T, wTxn txn.TxnMeta, s *service, commitTS int64, k byte, committed bool) {
   798  	for {
   799  		v := s.getTxnContext(wTxn.ID)
   800  		if v != nil {
   801  			time.Sleep(time.Millisecond * 100)
   802  			continue
   803  		}
   804  		break
   805  	}
   806  
   807  	kv := s.storage.(*mem.KVTxnStorage)
   808  
   809  	if committed {
   810  		kv.RLock()
   811  		v, ok := kv.GetCommittedKV().Get(GetTestKey(k), NewTestTimestamp(commitTS))
   812  		assert.True(t, ok)
   813  		assert.Equal(t, GetTestValue(k, wTxn), v)
   814  		kv.RUnlock()
   815  	} else {
   816  		kv.RLock()
   817  		n := 0
   818  		kv.GetCommittedKV().AscendRange(GetTestKey(k),
   819  			NewTestTimestamp(commitTS).Next(),
   820  			NewTestTimestamp(math.MaxInt64), func(_ []byte, _ timestamp.Timestamp) {
   821  				n++
   822  			})
   823  		assert.Equal(t, 0, n)
   824  		kv.RUnlock()
   825  	}
   826  
   827  	kv.RLock()
   828  	v, ok := kv.GetUncommittedKV().Get(GetTestKey(k))
   829  	assert.False(t, ok)
   830  	assert.Empty(t, v)
   831  	kv.RUnlock()
   832  
   833  	assert.Nil(t, kv.GetUncommittedTxn(wTxn.ID))
   834  }
   835  
   836  func addTestWaiter(t *testing.T, s *service, wTxn txn.TxnMeta, status txn.TxnStatus) *waiter {
   837  	txnCtx := s.getTxnContext(wTxn.ID)
   838  	assert.NotNil(t, txnCtx)
   839  	w := acquireWaiter()
   840  	assert.True(t, txnCtx.addWaiter(wTxn.ID, w, status))
   841  	return w
   842  }
   843  
   844  func checkWaiter(t *testing.T, w *waiter, expectStatus txn.TxnStatus) {
   845  	status, err := w.wait(context.Background())
   846  	assert.NoError(t, err)
   847  	assert.Equal(t, expectStatus, status)
   848  }