github.com/huiliang/nomad@v0.2.1-0.20151124023127-7a8b664699ff/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 }