add ctx and better handler

peers/e2e_test.go

@@ -3,6 +3,7 @@
 package peers
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -15,7 +16,7 @@ import (
 
 func TestE2E(t *testing.T) {
 	success := make(chan bool)
-	a := Peer{Handler: HandlerFunc(func(u *sticktable.EntryUpdate) {
+	a := Peer{Handler: HandlerFunc(func(_ context.Context, u *sticktable.EntryUpdate) {
 		log.Println(u)
 		success <- true
 	})}

peers/example/dump/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 
 	"github.com/dropmorepackets/haproxy-go/peers"
@@ -10,7 +11,7 @@ import (
 func main() {
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
 
-	err := peers.ListenAndServe(":21000", peers.HandlerFunc(func(u *sticktable.EntryUpdate) {
+	err := peers.ListenAndServe(":21000", peers.HandlerFunc(func(_ context.Context, u *sticktable.EntryUpdate) {
 		log.Println(u.String())
 	}))
 	if err != nil {

peers/example/prometheus-exporter/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"net/http"
 
@@ -24,7 +25,7 @@ func main() {
 
 	go http.ListenAndServe(":8081", promhttp.Handler())
 
-	err := peers.ListenAndServe(":21000", peers.HandlerFunc(func(update *sticktable.EntryUpdate) {
+	err := peers.ListenAndServe(":21000", peers.HandlerFunc(func(_ context.Context, update *sticktable.EntryUpdate) {
 		for i, d := range update.Data {
 			dt := update.StickTable.DataTypes[i].DataType
 			switch d := d.(type) {

peers/handler.go

@@ -1,13 +1,27 @@
 package peers
 
-import "github.com/dropmorepackets/haproxy-go/peers/sticktable"
+import (
+	"context"
+
+	"github.com/dropmorepackets/haproxy-go/peers/sticktable"
+)
 
 type Handler interface {
-	Update(*sticktable.EntryUpdate)
+	HandleUpdate(context.Context, *sticktable.EntryUpdate)
+	HandleHandshake(context.Context, *Handshake)
+	Close() error
 }
 
-type HandlerFunc func(*sticktable.EntryUpdate)
+type HandlerFunc func(context.Context, *sticktable.EntryUpdate)
+
+func (HandlerFunc) Close() error { return nil }
 
-func (h HandlerFunc) Update(u *sticktable.EntryUpdate) {
-	h(u)
+func (HandlerFunc) HandleHandshake(context.Context, *Handshake) {}
+
+func (h HandlerFunc) HandleUpdate(ctx context.Context, u *sticktable.EntryUpdate) {
+	h(ctx, u)
 }
+
+var (
+	_ Handler = (HandlerFunc)(nil)
+)

peers/peers.go

@@ -1,17 +1,17 @@
 package peers
 
 import (
-	"bufio"
 	"context"
 	"fmt"
 	"log"
 	"net"
 )
 
 type Peer struct {
-	Addr        string
-	Handler     Handler
-	BaseContext context.Context
+	Addr          string
+	Handler       Handler
+	HandlerSource func() Handler
+	BaseContext   context.Context
 }
 
 func ListenAndServe(addr string, handler Handler) error {
@@ -40,23 +40,28 @@ func (a *Peer) Serve(l net.Listener) error {
 		l.Close()
 	}()
 
+	if a.Handler != nil && a.HandlerSource != nil {
+		return fmt.Errorf("cannot set Handler and HandlerSource at the same time")
+	}
+
+	if a.Handler != nil {
+		a.HandlerSource = func() Handler {
+			return a.Handler
+		}
+	}
+
 	for {
 		nc, err := l.Accept()
 		if err != nil {
 			return fmt.Errorf("accepting conn: %w", err)
 		}
 
-		conn := &Conn{
-			ctx:     a.BaseContext,
-			conn:    nc,
-			r:       bufio.NewReader(nc),
-			handler: a.Handler,
-		}
-
+		p := newProtocolClient(a.BaseContext, nc, a.HandlerSource())
 		go func() {
 			defer nc.Close()
+			defer p.Close()
 
-			if err := conn.Serve(); err != nil && err != conn.ctx.Err() {
+			if err := p.Serve(); err != nil && err != p.ctx.Err() {
 				log.Println(err)
 			}
 		}()