github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/nomad/raft_rpc.go (about)

     1  package nomad
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/hashicorp/nomad/helper/tlsutil"
    10  	"github.com/hashicorp/raft"
    11  )
    12  
    13  // RaftLayer implements the raft.StreamLayer interface,
    14  // so that we can use a single RPC layer for Raft and Nomad
    15  type RaftLayer struct {
    16  	// Addr is the listener address to return
    17  	addr net.Addr
    18  
    19  	// connCh is used to accept connections
    20  	connCh chan net.Conn
    21  
    22  	// TLS wrapper
    23  	tlsWrap tlsutil.Wrapper
    24  
    25  	// Tracks if we are closed
    26  	closed    bool
    27  	closeCh   chan struct{}
    28  	closeLock sync.Mutex
    29  }
    30  
    31  // NewRaftLayer is used to initialize a new RaftLayer which can
    32  // be used as a StreamLayer for Raft. If a tlsConfig is provided,
    33  // then the connection will use TLS.
    34  func NewRaftLayer(addr net.Addr, tlsWrap tlsutil.Wrapper) *RaftLayer {
    35  	layer := &RaftLayer{
    36  		addr:    addr,
    37  		connCh:  make(chan net.Conn),
    38  		tlsWrap: tlsWrap,
    39  		closeCh: make(chan struct{}),
    40  	}
    41  	return layer
    42  }
    43  
    44  // Handoff is used to hand off a connection to the
    45  // RaftLayer. This allows it to be Accept()'ed
    46  func (l *RaftLayer) Handoff(c net.Conn) error {
    47  	select {
    48  	case l.connCh <- c:
    49  		return nil
    50  	case <-l.closeCh:
    51  		return fmt.Errorf("Raft RPC layer closed")
    52  	}
    53  }
    54  
    55  // Accept is used to return connection which are
    56  // dialed to be used with the Raft layer
    57  func (l *RaftLayer) Accept() (net.Conn, error) {
    58  	select {
    59  	case conn := <-l.connCh:
    60  		return conn, nil
    61  	case <-l.closeCh:
    62  		return nil, fmt.Errorf("Raft RPC layer closed")
    63  	}
    64  }
    65  
    66  // Close is used to stop listening for Raft connections
    67  func (l *RaftLayer) Close() error {
    68  	l.closeLock.Lock()
    69  	defer l.closeLock.Unlock()
    70  
    71  	if !l.closed {
    72  		l.closed = true
    73  		close(l.closeCh)
    74  	}
    75  	return nil
    76  }
    77  
    78  // Addr is used to return the address of the listener
    79  func (l *RaftLayer) Addr() net.Addr {
    80  	return l.addr
    81  }
    82  
    83  // Dial is used to create a new outgoing connection
    84  func (l *RaftLayer) Dial(address raft.ServerAddress, timeout time.Duration) (net.Conn, error) {
    85  	conn, err := net.DialTimeout("tcp", string(address), timeout)
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  
    90  	// Check for tls mode
    91  	if l.tlsWrap != nil {
    92  		// Switch the connection into TLS mode
    93  		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
    94  			conn.Close()
    95  			return nil, err
    96  		}
    97  
    98  		// Wrap the connection in a TLS client
    99  		conn, err = l.tlsWrap(conn)
   100  		if err != nil {
   101  			return nil, err
   102  		}
   103  	}
   104  
   105  	// Write the Raft byte to set the mode
   106  	_, err = conn.Write([]byte{byte(rpcRaft)})
   107  	if err != nil {
   108  		conn.Close()
   109  		return nil, err
   110  	}
   111  	return conn, err
   112  }