github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/closedts/transport/transport_util_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package transport_test
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/ctpb"
    17  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/transport"
    18  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/transport/testutils"
    19  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    20  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    21  	"github.com/cockroachdb/cockroach/pkg/util/stop"
    22  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    23  )
    24  
    25  type TestContainer struct {
    26  	Settings  *cluster.Settings
    27  	Stopper   *stop.Stopper
    28  	Producer  *TestProducer
    29  	Notifyee  *TestNotifyee
    30  	Refreshed *RefreshTracker
    31  	Server    *transport.Server
    32  	Dialer    *testutils.ChanDialer
    33  	Clients   *transport.Clients
    34  }
    35  
    36  type TestProducer struct {
    37  	syncutil.Mutex
    38  	chs []chan<- ctpb.Entry
    39  }
    40  
    41  func (tp *TestProducer) Subscribe(ctx context.Context, ch chan<- ctpb.Entry) {
    42  	tp.Lock()
    43  	tp.chs = append(tp.chs, ch)
    44  	tp.Unlock()
    45  }
    46  
    47  func (tp *TestProducer) numSubscriptions() int {
    48  	tp.Lock()
    49  	defer tp.Unlock()
    50  	return len(tp.chs)
    51  }
    52  
    53  func (tp *TestProducer) sendAll(entry ctpb.Entry) {
    54  	tp.Lock()
    55  	for _, ch := range tp.chs {
    56  		ch <- entry
    57  	}
    58  	tp.Unlock()
    59  }
    60  
    61  type TestNotifyee struct {
    62  	stopper *stop.Stopper
    63  	mu      struct {
    64  		syncutil.Mutex
    65  		entries map[roachpb.NodeID][]ctpb.Entry
    66  	}
    67  }
    68  
    69  func newTestNotifyee(stopper *stop.Stopper) *TestNotifyee {
    70  	tn := &TestNotifyee{
    71  		stopper: stopper,
    72  	}
    73  	tn.mu.entries = make(map[roachpb.NodeID][]ctpb.Entry)
    74  	return tn
    75  }
    76  
    77  func (tn *TestNotifyee) Notify(nodeID roachpb.NodeID) chan<- ctpb.Entry {
    78  	ch := make(chan ctpb.Entry)
    79  	tn.stopper.RunWorker(context.Background(), func(ctx context.Context) {
    80  		for entry := range ch {
    81  			tn.mu.Lock()
    82  			tn.mu.entries[nodeID] = append(tn.mu.entries[nodeID], entry)
    83  			tn.mu.Unlock()
    84  		}
    85  	})
    86  	return ch
    87  }
    88  
    89  type RefreshTracker struct {
    90  	syncutil.Mutex
    91  	rangeIDs []roachpb.RangeID
    92  }
    93  
    94  func (r *RefreshTracker) Get() []roachpb.RangeID {
    95  	r.Lock()
    96  	defer r.Unlock()
    97  	return append([]roachpb.RangeID(nil), r.rangeIDs...)
    98  }
    99  
   100  func (r *RefreshTracker) Add(rangeIDs ...roachpb.RangeID) {
   101  	r.Lock()
   102  	r.rangeIDs = append(r.rangeIDs, rangeIDs...)
   103  	r.Unlock()
   104  }