github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/causality/tests/driver.go (about)

     1  // Copyright 2022 PingCAP, 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 tests
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	"github.com/pingcap/errors"
    21  	"github.com/pingcap/log"
    22  	"github.com/pingcap/tiflow/pkg/causality"
    23  	"go.uber.org/atomic"
    24  	"go.uber.org/zap"
    25  )
    26  
    27  type conflictTestDriver struct {
    28  	workers          []*workerForTest
    29  	conflictDetector *causality.ConflictDetector[*txnForTest]
    30  	generator        workloadGenerator
    31  
    32  	pendingCount atomic.Int64
    33  }
    34  
    35  //nolint:unparam
    36  func newConflictTestDriver(
    37  	numWorkers int, numSlots int, workload workloadGenerator,
    38  ) *conflictTestDriver {
    39  	detector := causality.NewConflictDetector[*txnForTest](uint64(numSlots), causality.TxnCacheOption{
    40  		Count:         numWorkers,
    41  		Size:          1024,
    42  		BlockStrategy: causality.BlockStrategyWaitAvailable,
    43  	})
    44  
    45  	workers := make([]*workerForTest, 0, numWorkers)
    46  	for i := 0; i < numWorkers; i++ {
    47  		id := int64(i)
    48  		workers = append(workers, newWorkerForTest(detector.GetOutChByCacheID(id)))
    49  	}
    50  	return &conflictTestDriver{
    51  		workers:          workers,
    52  		conflictDetector: detector,
    53  		generator:        workload,
    54  	}
    55  }
    56  
    57  func (d *conflictTestDriver) WithExecFunc(fn func(txn *txnForTest) error) *conflictTestDriver {
    58  	for _, worker := range d.workers {
    59  		worker.execFunc = fn
    60  	}
    61  	return d
    62  }
    63  
    64  func (d *conflictTestDriver) Run(ctx context.Context, n int) error {
    65  	statusTicker := time.NewTicker(1 * time.Second)
    66  	defer statusTicker.Stop()
    67  
    68  	var counter int
    69  	for {
    70  		select {
    71  		case <-ctx.Done():
    72  			return errors.Trace(ctx.Err())
    73  		case <-statusTicker.C:
    74  			log.Info("Generation progress", zap.Int("count", counter))
    75  		default:
    76  		}
    77  
    78  		txn := &txnForTest{
    79  			keys: d.generator.Next(),
    80  			done: func() {
    81  				d.pendingCount.Sub(1)
    82  			},
    83  		}
    84  
    85  		d.pendingCount.Add(1)
    86  		d.conflictDetector.Add(txn)
    87  		counter++
    88  
    89  		if counter > n {
    90  			return nil
    91  		}
    92  	}
    93  }
    94  
    95  func (d *conflictTestDriver) Wait(ctx context.Context) error {
    96  	for {
    97  		select {
    98  		case <-ctx.Done():
    99  			return errors.Trace(ctx.Err())
   100  		default:
   101  		}
   102  
   103  		if d.pendingCount.Load() == 0 {
   104  			return nil
   105  		}
   106  
   107  		time.Sleep(10 * time.Millisecond)
   108  	}
   109  }
   110  
   111  func (d *conflictTestDriver) Close() {
   112  	d.conflictDetector.Close()
   113  	for _, worker := range d.workers {
   114  		worker.Close()
   115  	}
   116  }