Basic prototype

Author: abiosoft

tcpforward/balancer.go renamed to balancer/balancer.go

@@ -1,4 +1,4 @@
-package tcpforward
+package balancer
 
 import (
 	"encoding/json"
@@ -7,9 +7,11 @@ import (
 	"net"
 	"net/http"
 	"sync"
-	"github.com/abiosoft/dockward/util"
 )
 
+// TODO make dynamic
+const EndpointPort = 9923
+
 type Message struct {
 	Endpoint Endpoint
 	Remove   bool
@@ -22,6 +24,13 @@ type Balancer struct {
 	sync.RWMutex
 }
 
+func New(port int, endpoints Endpoints) *Balancer {
+	return &Balancer{
+		Port:      port,
+		Endpoints: endpoints,
+	}
+}
+
 func (b *Balancer) Start(stop chan struct{}) error {
 	listener, err := net.Listen("tcp", ":"+fmt.Sprint(b.Port))
 	if err != nil {
@@ -71,36 +80,29 @@ func (b *Balancer) Select(e Endpoints) Endpoint {
 	return b.Policy.Select(e)
 }
 
-func (b *Balancer) ListenForEndpoints() (int, error) {
-	port, err := util.RandomPort()
-	if err != nil {
-		return port, err
-	}
+func (b *Balancer) ListenForEndpoints(port int) {
+	handler := http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			var message Message
+			err := json.NewDecoder(r.Body).Decode(&message)
+			if r.Method != "POST" || err != nil {
+				http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
+				return
+			}
 
-	go func() {
-		err := http.ListenAndServe(":"+fmt.Sprint(port),
-			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				var message Message
-				err := json.NewDecoder(r.Body).Decode(&message)
-				if err != nil {
-					http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
-					return
-				}
-
-				b.Lock()
-				if message.Remove {
-					b.Endpoints.Delete(message.Endpoint.Id)
-				} else {
-					b.Endpoints.Add(message.Endpoint)
-				}
-				b.Unlock()
-
-				w.WriteHeader(200)
-			}))
-
-		// should not get here
-		log.Println(err)
-	}()
+			b.Lock()
+			if message.Remove {
+				b.Endpoints.Delete(message.Endpoint.Ip)
+			} else {
+				b.Endpoints.Add(message.Endpoint)
+			}
+			b.Unlock()
+
+			w.WriteHeader(200)
+		})
+
+	err := http.ListenAndServe(":"+fmt.Sprint(port), handler)
 
-	return port, err
+	// should not get here
+	log.Println(err)
 }

balancer/endpoint.go

@@ -0,0 +1,88 @@
+package balancer
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/abiosoft/dockward/util"
+)
+
+type Endpoint struct {
+	Id   string
+	Ip   string
+	Port int
+}
+
+func ParseEndpoint(addr string) Endpoint {
+	// assume addr as host, port as 80
+	ip, port, id := addr, 80, util.RandChars(10)
+
+	// if its valid port, assume as port, host as 127.0.0.1
+	if p, err := strconv.Atoi(addr); err == nil {
+		ip = "127.0.0.1"
+		port = p
+	}
+
+	// attempt parse
+	str := strings.Split(addr, ":")
+
+	// valid host/port
+	if len(str) > 1 {
+		ip = str[0]
+		port, _ = strconv.Atoi(str[1])
+	}
+	// valid id
+	if len(str) > 2 {
+		id = str[2]
+	}
+
+	return Endpoint{
+		Id:   id,
+		Ip:   ip,
+		Port: port,
+	}
+}
+
+func (ep Endpoint) Addr() string {
+	return ep.Ip + ":" + fmt.Sprint(ep.Port)
+}
+
+func (ep Endpoint) String() string {
+	return ep.Addr() + ":" + ep.Id
+}
+
+type Endpoints []Endpoint
+
+func (e Endpoints) Len() int {
+	return len(e)
+}
+
+func (e *Endpoints) Add(ep Endpoint) {
+	for i, endpoint := range *e {
+		if endpoint.Id == ep.Id {
+			// already exists, replace instead.
+			(*e)[i] = ep
+			return
+		}
+	}
+	*e = append(*e, ep)
+}
+
+func (e *Endpoints) Delete(id string) {
+	pos := -1
+	for i, ep := range *e {
+		if ep.Id == id {
+			pos = i
+			break
+		}
+	}
+	if pos == -1 {
+		return
+	}
+	part := (*e)[:pos]
+	if pos < len(*e)-1 {
+		part = append(part, (*e)[pos+1:]...)
+	}
+	*e = part
+}

tcpforward/schedule.go renamed to balancer/schedule.go

@@ -1,6 +1,9 @@
-package tcpforward
+package balancer
 
-import "math/rand"
+import (
+	"math/rand"
+	"time"
+)
 
 // Policy is a selection policy.
 type Policy interface {
@@ -12,9 +15,12 @@ type Random struct{}
 
 // Select satisfies Policy.
 func (r Random) Select(d Endpoints) Endpoint {
-	i := rand.Int() % (len(d) - 1)
-	if i < 0 {
-		i = 0
+	if len(d) == 0 {
+		return Endpoint{}
 	}
-	return d[i]
+	return d[rand.Int()%len(d)]
+}
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
 }

tcpforward/tcpforward.go renamed to balancer/tcpforward.go

@@ -1,4 +1,4 @@
-package tcpforward
+package balancer
 
 import (
 	"fmt"
@@ -32,6 +32,9 @@ func handle(conn net.Conn, dest string) {
 	client, err := net.Dial("tcp", dest)
 	if err != nil {
 		log.Println(err)
+		if err := conn.Close(); err != nil {
+			log.Println(err)
+		}
 		return
 	}
 	var w sync.WaitGroup

cleanup.go

@@ -0,0 +1,44 @@
+package main
+import (
+	"os/signal"
+	"fmt"
+	"syscall"
+	"os"
+)
+
+// trapInterrupts traps OS interrupt signals.
+func trapInterrupts(exit chan struct{}) chan struct{} {
+	done := make(chan struct{})
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+	signal.Notify(c, syscall.SIGTERM)
+	go func() {
+		select {
+		case <-c:
+			fmt.Print("OS Interrupt signal received. Performing cleanup...")
+			cleanUp()
+			fmt.Println("Done.")
+			done <- struct{}{}
+		case <-exit:
+			cleanUp()
+			done <- struct{}{}
+		}
+
+	}()
+	return done
+}
+
+// cleanUpFuncs is list of functions to call before application exits.
+var cleanUpFuncs []func()
+
+// addCleanUpFunc adds a function to cleanUpFuncs.
+func addCleanUpFunc(f func()) {
+	cleanUpFuncs = append(cleanUpFuncs, f)
+}
+
+// cleanUp calls all functions in cleanUpFuncs.
+func cleanUp() {
+	for i := range cleanUpFuncs {
+		cleanUpFuncs[i]()
+	}
+}

cli.go

@@ -0,0 +1,69 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"strconv"
+)
+
+type cliArgs struct {
+	HostPort       int
+	ContainerName  string
+	ContainerId    string
+	ContainerLabel string
+	Host           bool
+	Endpoints      []string
+}
+
+func usageErr(err error) {
+	exit(fmt.Errorf("%v\n\n%v", err, Usage))
+}
+
+func parseCli() cliArgs {
+	if len(os.Args) == 1 {
+		usageErr(fmt.Errorf("Command missing"))
+	}
+
+	switch os.Args[1] {
+	case "help":
+		fmt.Println(Usage)
+		exit(nil)
+	case "version":
+		fmt.Println("dockward version", Version)
+		exit(nil)
+	}
+	hostPort, err := strconv.Atoi(os.Args[1])
+	if err != nil {
+		usageErr(err)
+	}
+
+	args := cliArgs{HostPort: hostPort}
+
+	fs := flag.FlagSet{}
+	fs.SetOutput(ioutil.Discard)
+
+	fs.BoolVar(&args.Host, "host", args.Host, "")
+	fs.StringVar(&args.ContainerId, "id", args.ContainerId, "")
+	fs.StringVar(&args.ContainerName, "name", args.ContainerName, "")
+	fs.StringVar(&args.ContainerLabel, "label", args.ContainerLabel, "")
+
+	err = fs.Parse(os.Args[2:])
+	if err != nil {
+		exit(err)
+	}
+
+	// if not host mode, require a container param.
+	if !args.Host {
+		if args.ContainerId == "" && args.ContainerLabel == "" && args.ContainerName == "" {
+			exit(fmt.Errorf("One of container id, name or label is required."))
+		}
+	}
+
+	if fs.NArg() > 0 {
+		args.Endpoints = fs.Args()
+	}
+	return args
+}
+