go.uber.org/yarpc@v1.72.1/peer/hostport/hostport.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 hostport
    22  
    23  import (
    24  	"sync"
    25  
    26  	"go.uber.org/atomic"
    27  	"go.uber.org/yarpc/api/peer"
    28  )
    29  
    30  // PeerIdentifier uniquely references a host:port combination using a common interface
    31  type PeerIdentifier string
    32  
    33  // Identifier generates a (should be) unique identifier for this PeerIdentifier (to use in maps, etc)
    34  func (p PeerIdentifier) Identifier() string {
    35  	return string(p)
    36  }
    37  
    38  // Identify coerces a string to a PeerIdentifier
    39  func Identify(peer string) peer.Identifier {
    40  	return PeerIdentifier(peer)
    41  }
    42  
    43  // NewPeer creates a new hostport.Peer from a hostport.PeerIdentifier, peer.Transport, and peer.Subscriber
    44  func NewPeer(pid PeerIdentifier, transport peer.Transport) *Peer {
    45  	p := &Peer{
    46  		PeerIdentifier: pid,
    47  		transport:      transport,
    48  		subscribers:    make(map[peer.Subscriber]struct{}),
    49  	}
    50  	p.connectionStatus.Store(int32(peer.Unavailable))
    51  	return p
    52  }
    53  
    54  // Peer keeps a subscriber to send status updates to it, and the peer.Transport that created it
    55  type Peer struct {
    56  	PeerIdentifier
    57  
    58  	lock sync.RWMutex
    59  
    60  	transport        peer.Transport
    61  	subscribers      map[peer.Subscriber]struct{}
    62  	pending          atomic.Int32
    63  	connectionStatus atomic.Int32
    64  }
    65  
    66  // HostPort surfaces the HostPort in this function, if you want to access the hostport directly (for a downstream call)
    67  // You need to cast the Peer to a *hostport.Peer and run this function
    68  func (p *Peer) HostPort() string {
    69  	return string(p.PeerIdentifier)
    70  }
    71  
    72  // Transport returns the peer.Transport that is in charge of this hostport.Peer (and should be the one to handle requests)
    73  func (p *Peer) Transport() peer.Transport {
    74  	return p.transport
    75  }
    76  
    77  // Subscribe adds a subscriber to the peer's subscriber map
    78  func (p *Peer) Subscribe(sub peer.Subscriber) {
    79  	p.lock.Lock()
    80  	p.subscribers[sub] = struct{}{}
    81  	p.lock.Unlock()
    82  }
    83  
    84  // Unsubscribe removes a subscriber from the peer's subscriber map
    85  func (p *Peer) Unsubscribe(sub peer.Subscriber) error {
    86  	p.lock.Lock()
    87  	defer p.lock.Unlock()
    88  	if _, ok := p.subscribers[sub]; !ok {
    89  		return peer.ErrPeerHasNoReferenceToSubscriber{
    90  			PeerIdentifier: p.PeerIdentifier,
    91  			PeerSubscriber: sub,
    92  		}
    93  	}
    94  
    95  	delete(p.subscribers, sub)
    96  	return nil
    97  }
    98  
    99  // NumSubscribers returns the number of subscriptions attached to the peer
   100  func (p *Peer) NumSubscribers() int {
   101  	p.lock.RLock()
   102  	subs := len(p.subscribers)
   103  	p.lock.RUnlock()
   104  	return subs
   105  }
   106  
   107  // Status returns the current status of the hostport.Peer
   108  func (p *Peer) Status() peer.Status {
   109  	return peer.Status{
   110  		PendingRequestCount: int(p.pending.Load()),
   111  		ConnectionStatus:    peer.ConnectionStatus(p.connectionStatus.Load()),
   112  	}
   113  }
   114  
   115  // SetStatus sets the status of the Peer (to be used by the peer.Transport)
   116  func (p *Peer) SetStatus(status peer.ConnectionStatus) {
   117  	p.connectionStatus.Store(int32(status))
   118  	p.notifyStatusChanged()
   119  }
   120  
   121  // StartRequest runs at the beginning of a request.
   122  func (p *Peer) StartRequest() {
   123  	p.pending.Inc()
   124  	p.notifyStatusChanged()
   125  }
   126  
   127  // EndRequest should be run after a request has finished.
   128  func (p *Peer) EndRequest() {
   129  	p.pending.Dec()
   130  	p.notifyStatusChanged()
   131  }
   132  
   133  func (p *Peer) notifyStatusChanged() {
   134  	p.lock.RLock()
   135  	subs := make([]peer.Subscriber, 0, len(p.subscribers))
   136  	for sub := range p.subscribers {
   137  		subs = append(subs, sub)
   138  	}
   139  	p.lock.RUnlock()
   140  
   141  	for _, sub := range subs {
   142  		sub.NotifyStatusChanged(p)
   143  	}
   144  }