github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/closedts/transport/testutils/chan_dialer.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 testutils
    12  
    13  import (
    14  	"context"
    15  	"io"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/ctpb"
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/util/stop"
    20  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    21  )
    22  
    23  // ChanDialer is an implementation of closedts.Dialer that connects clients
    24  // directly via a channel to a Server.
    25  type ChanDialer struct {
    26  	stopper *stop.Stopper
    27  	server  ctpb.Server
    28  
    29  	mu struct {
    30  		syncutil.Mutex
    31  		transcripts map[roachpb.NodeID][]interface{}
    32  	}
    33  }
    34  
    35  // NewChanDialer sets up a ChanDialer.
    36  func NewChanDialer(stopper *stop.Stopper, server ctpb.Server) *ChanDialer {
    37  	d := &ChanDialer{
    38  		stopper: stopper,
    39  		server:  server,
    40  	}
    41  	d.mu.transcripts = make(map[roachpb.NodeID][]interface{})
    42  	return d
    43  }
    44  
    45  // Transcript returns a slice of messages sent over the "wire".
    46  func (d *ChanDialer) Transcript(nodeID roachpb.NodeID) []interface{} {
    47  	d.mu.Lock()
    48  	defer d.mu.Unlock()
    49  	return append([]interface{}(nil), d.mu.transcripts[nodeID]...)
    50  }
    51  
    52  // Dial implements closedts.Dialer.
    53  func (d *ChanDialer) Dial(ctx context.Context, nodeID roachpb.NodeID) (ctpb.Client, error) {
    54  	c := &client{
    55  		ctx:     ctx,
    56  		send:    make(chan *ctpb.Reaction),
    57  		recv:    make(chan *ctpb.Entry),
    58  		stopper: d.stopper,
    59  		observe: func(msg interface{}) {
    60  			d.mu.Lock()
    61  			if d.mu.transcripts == nil {
    62  				d.mu.transcripts = map[roachpb.NodeID][]interface{}{}
    63  			}
    64  			d.mu.transcripts[nodeID] = append(d.mu.transcripts[nodeID], msg)
    65  			d.mu.Unlock()
    66  		},
    67  	}
    68  
    69  	d.stopper.RunWorker(ctx, func(ctx context.Context) {
    70  		_ = d.server.Get((*incomingClient)(c))
    71  	})
    72  	return c, nil
    73  
    74  }
    75  
    76  // Ready implements closedts.Dialer by always returning true.
    77  func (d *ChanDialer) Ready(nodeID roachpb.NodeID) bool {
    78  	return true
    79  }
    80  
    81  type client struct {
    82  	ctx     context.Context
    83  	stopper *stop.Stopper
    84  	send    chan *ctpb.Reaction
    85  	recv    chan *ctpb.Entry
    86  
    87  	observe func(interface{})
    88  }
    89  
    90  func (c *client) Send(msg *ctpb.Reaction) error {
    91  	select {
    92  	case <-c.stopper.ShouldQuiesce():
    93  		return io.EOF
    94  	case c.send <- msg:
    95  		c.observe(msg)
    96  		return nil
    97  	}
    98  }
    99  
   100  func (c *client) Recv() (*ctpb.Entry, error) {
   101  	select {
   102  	case <-c.stopper.ShouldQuiesce():
   103  		return nil, io.EOF
   104  	case msg := <-c.recv:
   105  		c.observe(msg)
   106  		return msg, nil
   107  	}
   108  }
   109  
   110  func (c *client) CloseSend() error {
   111  	close(c.send)
   112  	return nil
   113  }
   114  
   115  func (c *client) Context() context.Context {
   116  	return c.ctx
   117  }
   118  
   119  type incomingClient client
   120  
   121  func (c *incomingClient) Send(msg *ctpb.Entry) error {
   122  	select {
   123  	case <-c.stopper.ShouldQuiesce():
   124  		return io.EOF
   125  	case c.recv <- msg:
   126  		return nil
   127  	}
   128  }
   129  
   130  func (c *incomingClient) Recv() (*ctpb.Reaction, error) {
   131  	select {
   132  	case <-c.stopper.ShouldQuiesce():
   133  		return nil, io.EOF
   134  	case msg, ok := <-c.send:
   135  		if !ok {
   136  			return nil, io.EOF
   137  		}
   138  		return msg, nil
   139  	}
   140  }
   141  
   142  func (c *incomingClient) Context() context.Context {
   143  	return c.ctx
   144  }