github.com/hashicorp/go-plugin@v1.6.0/internal/grpcmux/grpc_client_muxer.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package grpcmux
     5  
     6  import (
     7  	"fmt"
     8  	"net"
     9  	"sync"
    10  
    11  	"github.com/hashicorp/go-hclog"
    12  	"github.com/hashicorp/yamux"
    13  )
    14  
    15  var _ GRPCMuxer = (*GRPCClientMuxer)(nil)
    16  
    17  // GRPCClientMuxer implements the client (host) side of the gRPC broker's
    18  // GRPCMuxer interface for multiplexing multiple gRPC broker connections over
    19  // a single net.Conn.
    20  //
    21  // The client dials the initial net.Conn eagerly, and creates a yamux.Session
    22  // as the implementation for multiplexing any additional connections.
    23  //
    24  // Each net.Listener returned from Listener will block until the client receives
    25  // a knock that matches its gRPC broker stream ID. There is no default listener
    26  // on the client, as it is a client for the gRPC broker's control services. (See
    27  // GRPCServerMuxer for more details).
    28  type GRPCClientMuxer struct {
    29  	logger  hclog.Logger
    30  	session *yamux.Session
    31  
    32  	acceptMutex     sync.Mutex
    33  	acceptListeners map[uint32]*blockedClientListener
    34  }
    35  
    36  func NewGRPCClientMuxer(logger hclog.Logger, addr net.Addr) (*GRPCClientMuxer, error) {
    37  	// Eagerly establish the underlying connection as early as possible.
    38  	logger.Debug("making new client mux initial connection", "addr", addr)
    39  	conn, err := net.Dial(addr.Network(), addr.String())
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	if tcpConn, ok := conn.(*net.TCPConn); ok {
    44  		// Make sure to set keep alive so that the connection doesn't die
    45  		_ = tcpConn.SetKeepAlive(true)
    46  	}
    47  
    48  	cfg := yamux.DefaultConfig()
    49  	cfg.Logger = logger.Named("yamux").StandardLogger(&hclog.StandardLoggerOptions{
    50  		InferLevels: true,
    51  	})
    52  	cfg.LogOutput = nil
    53  	sess, err := yamux.Client(conn, cfg)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  
    58  	logger.Debug("client muxer connected", "addr", addr)
    59  	m := &GRPCClientMuxer{
    60  		logger:          logger,
    61  		session:         sess,
    62  		acceptListeners: make(map[uint32]*blockedClientListener),
    63  	}
    64  
    65  	return m, nil
    66  }
    67  
    68  func (m *GRPCClientMuxer) Enabled() bool {
    69  	return m != nil
    70  }
    71  
    72  func (m *GRPCClientMuxer) Listener(id uint32, doneCh <-chan struct{}) (net.Listener, error) {
    73  	ln := newBlockedClientListener(m.session, doneCh)
    74  
    75  	m.acceptMutex.Lock()
    76  	m.acceptListeners[id] = ln
    77  	m.acceptMutex.Unlock()
    78  
    79  	return ln, nil
    80  }
    81  
    82  func (m *GRPCClientMuxer) AcceptKnock(id uint32) error {
    83  	m.acceptMutex.Lock()
    84  	defer m.acceptMutex.Unlock()
    85  
    86  	ln, ok := m.acceptListeners[id]
    87  	if !ok {
    88  		return fmt.Errorf("no listener for id %d", id)
    89  	}
    90  	ln.unblock()
    91  	return nil
    92  }
    93  
    94  func (m *GRPCClientMuxer) Dial() (net.Conn, error) {
    95  	stream, err := m.session.Open()
    96  	if err != nil {
    97  		return nil, fmt.Errorf("error dialling new client stream: %w", err)
    98  	}
    99  
   100  	return stream, nil
   101  }
   102  
   103  func (m *GRPCClientMuxer) Close() error {
   104  	return m.session.Close()
   105  }