Refactor proxy code to make room for UDP

pkg/proxy/proxier.go

@@ -21,6 +21,7 @@ import (
 	"io"
 	"net"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -32,11 +33,89 @@ import (
 type serviceInfo struct {
 	name     string
 	port     int
-	listener net.Listener
+	protocol string
+	socket   proxySocket
 	mu       sync.Mutex // protects active
 	active   bool
 }
 
+// Abstraction over TCP/UDP sockets which are proxied.
+type proxySocket interface {
+	// Addr gets the net.Addr for a proxySocket.
+	Addr() net.Addr
+	// Close stops the proxySocket from accepting incoming connections.
+	Close() error
+	// ProxyLoop proxies incoming connections for the specified service to the service endpoints.
+	ProxyLoop(service string, proxier *Proxier)
+}
+
+type tcpProxySocket struct {
+	net.Listener
+}
+
+func (tcp *tcpProxySocket) ProxyLoop(service string, proxier *Proxier) {
+	info, found := proxier.getServiceInfo(service)
+	if !found {
+		glog.Errorf("Failed to find service: %s", service)
+		return
+	}
+	for {
+		info.mu.Lock()
+		if !info.active {
+			info.mu.Unlock()
+			break
+		}
+		info.mu.Unlock()
+
+		// Block until a connection is made.
+		inConn, err := tcp.Accept()
+		if err != nil {
+			glog.Errorf("Accept failed: %v", err)
+			continue
+		}
+		glog.Infof("Accepted connection from %v to %v", inConn.RemoteAddr(), inConn.LocalAddr())
+		endpoint, err := proxier.loadBalancer.NextEndpoint(service, inConn.RemoteAddr())
+		if err != nil {
+			glog.Errorf("Couldn't find an endpoint for %s %v", service, err)
+			inConn.Close()
+			continue
+		}
+		glog.Infof("Mapped service %s to endpoint %s", service, endpoint)
+		// TODO: This could spin up a new goroutine to make the outbound connection,
+		// and keep accepting inbound traffic.
+		outConn, err := net.DialTimeout("tcp", endpoint, time.Duration(5)*time.Second)
+		if err != nil {
+			// TODO: Try another endpoint?
+			glog.Errorf("Dial failed: %v", err)
+			inConn.Close()
+			continue
+		}
+		// Spin up an async copy loop.
+		proxyTCP(inConn.(*net.TCPConn), outConn.(*net.TCPConn))
+	}
+}
+
+// proxyTCP proxies data bi-directionally between in and out.
+func proxyTCP(in, out *net.TCPConn) {
+	glog.Infof("Creating proxy between %v <-> %v <-> %v <-> %v",
+		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
+	go copyBytes(in, out)
+	go copyBytes(out, in)
+}
+
+func newProxySocket(protocol string, addr string, port int) (proxySocket, error) {
+	switch strings.ToUpper(protocol) {
+	case "TCP":
+		listener, err := net.Listen("tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
+		if err != nil {
+			return nil, err
+		}
+		return &tcpProxySocket{listener}, nil
+		//TODO: add UDP support
+	}
+	return nil, fmt.Errorf("Unknown protocol %q", protocol)
+}
+
 // Proxier is a simple proxy for TCP connections between a localhost:lport
 // and services that provide the actual implementations.
 type Proxier struct {
@@ -66,14 +145,6 @@ func copyBytes(in, out *net.TCPConn) {
 	out.CloseWrite()
 }
 
-// proxyConnection proxies data bidirectionally between in and out.
-func proxyConnection(in, out *net.TCPConn) {
-	glog.Infof("Creating proxy between %v <-> %v <-> %v <-> %v",
-		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
-	go copyBytes(in, out)
-	go copyBytes(out, in)
-}
-
 // StopProxy stops the proxy for the named service.
 func (proxier *Proxier) StopProxy(service string) error {
 	// TODO: delete from map here?
@@ -92,7 +163,7 @@ func (proxier *Proxier) stopProxyInternal(info *serviceInfo) error {
 	}
 	glog.Infof("Removing service: %s", info.name)
 	info.active = false
-	return info.listener.Close()
+	return info.socket.Close()
 }
 
 func (proxier *Proxier) getServiceInfo(service string) (*serviceInfo, bool) {
@@ -109,57 +180,19 @@ func (proxier *Proxier) setServiceInfo(service string, info *serviceInfo) {
 	proxier.serviceMap[service] = info
 }
 
-// AcceptHandler proxies incoming connections for the specified service
-// to the load-balanced service endpoints.
-func (proxier *Proxier) AcceptHandler(service string, listener net.Listener) {
-	info, found := proxier.getServiceInfo(service)
-	if !found {
-		glog.Errorf("Failed to find service: %s", service)
-		return
-	}
-	for {
-		info.mu.Lock()
-		if !info.active {
-			info.mu.Unlock()
-			break
-		}
-		info.mu.Unlock()
-		inConn, err := listener.Accept()
-		if err != nil {
-			glog.Errorf("Accept failed: %v", err)
-			continue
-		}
-		glog.Infof("Accepted connection from: %v to %v", inConn.RemoteAddr(), inConn.LocalAddr())
-		endpoint, err := proxier.loadBalancer.NextEndpoint(service, inConn.RemoteAddr())
-		if err != nil {
-			glog.Errorf("Couldn't find an endpoint for %s %v", service, err)
-			inConn.Close()
-			continue
-		}
-		glog.Infof("Mapped service %s to endpoint %s", service, endpoint)
-		outConn, err := net.DialTimeout("tcp", endpoint, time.Duration(5)*time.Second)
-		if err != nil {
-			glog.Errorf("Dial failed: %v", err)
-			inConn.Close()
-			continue
-		}
-		proxyConnection(inConn.(*net.TCPConn), outConn.(*net.TCPConn))
-	}
-}
-
 // used to globally lock around unused ports. Only used in testing.
 var unusedPortLock sync.Mutex
 
 // addServiceOnUnusedPort starts listening for a new service, returning the
 // port it's using.  For testing on a system with unknown ports used.
-func (proxier *Proxier) addServiceOnUnusedPort(service string) (string, error) {
+func (proxier *Proxier) addServiceOnUnusedPort(service, protocol string) (string, error) {
 	unusedPortLock.Lock()
 	defer unusedPortLock.Unlock()
-	l, err := net.Listen("tcp", net.JoinHostPort(proxier.address, "0"))
+	sock, err := newProxySocket(protocol, proxier.address, 0)
 	if err != nil {
 		return "", err
 	}
-	_, port, err := net.SplitHostPort(l.Addr().String())
+	_, port, err := net.SplitHostPort(sock.Addr().String())
 	if err != nil {
 		return "", err
 	}
@@ -169,16 +202,17 @@ func (proxier *Proxier) addServiceOnUnusedPort(service string) (string, error) {
 	}
 	proxier.setServiceInfo(service, &serviceInfo{
 		port:     portNum,
+		protocol: protocol,
 		active:   true,
-		listener: l,
+		socket:   sock,
 	})
-	proxier.startAccepting(service, l)
+	proxier.startAccepting(service, sock)
 	return port, nil
 }
 
-func (proxier *Proxier) startAccepting(service string, l net.Listener) {
-	glog.Infof("Listening for %s on %s", service, l.Addr().String())
-	go proxier.AcceptHandler(service, l)
+func (proxier *Proxier) startAccepting(service string, sock proxySocket) {
+	glog.Infof("Listening for %s on %s", service, sock.Addr().String())
+	go sock.ProxyLoop(service, proxier)
 }
 
 // OnUpdate manages the active set of service proxies.
@@ -196,18 +230,19 @@ func (proxier *Proxier) OnUpdate(services []api.Service) {
 		if exists && info.port != service.Port {
 			proxier.StopProxy(service.ID)
 		}
-		glog.Infof("Adding a new service %s on port %d", service.ID, service.Port)
-		listener, err := net.Listen("tcp", net.JoinHostPort(proxier.address, strconv.Itoa(service.Port)))
+		glog.Infof("Adding a new service %s on %s port %d", service.ID, service.Protocol, service.Port)
+		sock, err := newProxySocket(service.Protocol, proxier.address, service.Port)
 		if err != nil {
-			glog.Infof("Failed to start listening for %s on %d", service.ID, service.Port)
+			glog.Errorf("Failed to get a socket for %s: %+v", service.ID, err)
 			continue
 		}
 		proxier.setServiceInfo(service.ID, &serviceInfo{
 			port:     service.Port,
+			protocol: service.Protocol,
 			active:   true,
-			listener: listener,
+			socket:   sock,
 		})
-		proxier.startAccepting(service.ID, listener)
+		proxier.startAccepting(service.ID, sock)
 	}
 	proxier.mu.Lock()
 	defer proxier.mu.Unlock()