go.uber.org/yarpc@v1.72.1/peer/direct/direct.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package direct
    22  
    23  import (
    24  	"context"
    25  	"fmt"
    26  
    27  	"go.uber.org/yarpc/api/peer"
    28  	"go.uber.org/yarpc/api/transport"
    29  	"go.uber.org/yarpc/peer/hostport"
    30  	"go.uber.org/yarpc/yarpcerrors"
    31  	"go.uber.org/zap"
    32  )
    33  
    34  var _ peer.Chooser = (*Chooser)(nil)
    35  
    36  type chooserOptions struct {
    37  	logger *zap.Logger
    38  }
    39  
    40  // ChooserOption customizes the behavior of the peer chooser.
    41  type ChooserOption func(*chooserOptions)
    42  
    43  // Logger sets the logger for the chooser.
    44  func Logger(logger *zap.Logger) ChooserOption {
    45  	return func(c *chooserOptions) {
    46  		c.logger = logger
    47  	}
    48  }
    49  
    50  // New creates a new direct peer chooser.
    51  func New(cfg Configuration, transport peer.Transport, opts ...ChooserOption) (*Chooser, error) {
    52  	options := &chooserOptions{
    53  		logger: zap.NewNop(),
    54  	}
    55  	for _, opt := range opts {
    56  		opt(options)
    57  	}
    58  
    59  	if transport == nil {
    60  		return nil, yarpcerrors.InvalidArgumentErrorf("%q chooser requires non-nil transport", name)
    61  	}
    62  	return &Chooser{
    63  		transport: transport,
    64  		logger:    options.logger,
    65  	}, nil
    66  }
    67  
    68  // Chooser is a peer.Chooser that returns the peer identified by
    69  // transport.Request#ShardKey, suitable for directly addressing a peer.
    70  type Chooser struct {
    71  	transport peer.Transport
    72  	logger    *zap.Logger
    73  }
    74  
    75  // Start statisfies the peer.Chooser interface.
    76  func (c *Chooser) Start() error {
    77  	return nil // no-op
    78  }
    79  
    80  // Stop statisfies the peer.Chooser interface.
    81  func (c *Chooser) Stop() error {
    82  	return nil // no-op
    83  }
    84  
    85  // IsRunning statisfies the peer.Chooser interface.
    86  func (c *Chooser) IsRunning() bool {
    87  	return true // no-op
    88  }
    89  
    90  // Choose uses the peer identifier set as the transport.Request#ShardKey to
    91  // return the peer.
    92  func (c *Chooser) Choose(ctx context.Context, req *transport.Request) (peer.Peer, func(error), error) {
    93  	if req.ShardKey == "" {
    94  		return nil, nil,
    95  			yarpcerrors.InvalidArgumentErrorf("%q chooser requires ShardKey to be non-empty", name)
    96  	}
    97  
    98  	id := hostport.Identify(req.ShardKey)
    99  	sub := &peerSubscriber{
   100  		peerIdentifier: id,
   101  	}
   102  
   103  	transportPeer, err := c.transport.RetainPeer(id, sub)
   104  	if err != nil {
   105  		return nil, nil, err
   106  	}
   107  
   108  	onFinish := func(error) {
   109  		if err := c.transport.ReleasePeer(transportPeer, sub); err != nil {
   110  			c.logger.Error(
   111  				fmt.Sprintf("error releasing peer from transport in %q chooser", name),
   112  				zap.Error(err))
   113  		}
   114  	}
   115  	return transportPeer, onFinish, nil
   116  }
   117  
   118  type peerSubscriber struct {
   119  	// Even if the peerIdentifier of peerSubscriber is not used,
   120  	// struct with no fields does not behave the same way as struct with fields.
   121  	// For instance, with no fields and p1 := &peerSubscriber{}, p2 := &peerSubscriber{}
   122  	// &p1 == &p2 will be true.
   123  	// Internally, YARPC stores this *peerSubscriber as a hash's key. p1 and p2 must be different.
   124  	// More details here: https://dave.cheney.net/2014/03/25/the-empty-struct
   125  	peerIdentifier peer.Identifier
   126  }
   127  
   128  func (d *peerSubscriber) NotifyStatusChanged(_ peer.Identifier) {}