github.com/kardianos/nomad@v0.1.3-0.20151022182107-b13df73ee850/nomad/raft_rpc.go (about)

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