feat: add UDP support (#20)

Author: antiphp

README.md

@@ -49,14 +49,14 @@ The `server` command supports the following additional arguments.
 
 ```shell
 OPTIONS:
-   --addr value           The address to listen on for probe messages (default: ":8123") [$ADDR]
-   --buffer-size value    The size of the read buffer used by the server (default: 512) [$BUFFER_SIZE]
-   --read-timeout value   The duration after which the server should timeout when reading from a connection (default: 2s) [$READ_TIMEOUT]
-   --write-timeout value  The duration after which the server should timeout when writing to a connection (default: 5s) [$WRITE_TIMEOUT]
-   --log.format value     Specify the format of logs. Supported formats: 'logfmt', 'json', 'console' [$LOG_FORMAT]
-   --log.level value      Specify the log level. e.g. 'debug', 'info', 'error'. (default: "info") [$LOG_LEVEL]
-   --log.ctx value        A list of context field appended to every log. Format: key=value. [$LOG_CTX]
-   --help, -h             show help (default: false)
+   --addr value                         The address to listen on for probe messages (default: ":8123") [$ADDR]
+   --buffer-size value                  The size of the read buffer used by the server (default: 512) [$BUFFER_SIZE]
+   --read-timeout value                 The duration after which the server should timeout when reading from a connection (default: 2s) [$READ_TIMEOUT]
+   --write-timeout value                The duration after which the server should timeout when writing to a connection (default: 5s) [$WRITE_TIMEOUT]
+   --log.format value                   Specify the format of logs. Supported formats: 'logfmt', 'json', 'console' [$LOG_FORMAT]
+   --log.level value                    Specify the log level. e.g. 'debug', 'info', 'error'. (default: "info") [$LOG_LEVEL]
+   --log.ctx value [ --log.ctx value ]  A list of context field appended to every log. Format: key=value. [$LOG_CTX]
+   --help, -h                           show help
 ```
 
 #### Firewall
@@ -95,15 +95,17 @@ The `client` command supports the following additional arguments.
 
 ```shell
 OPTIONS:
-   --addr value           The address of a connqc server [$ADDR]
-   --backoff value        Duration to wait for before retrying to connect to the server (default: 1s) [$BACKOFF]
-   --interval value       Interval at which to send probe messages to the server (default: 1s) [$INTERVAL]
-   --read-timeout value   The duration after which the client should timeout when reading from a connection (default: 2s) [$READ_TIMEOUT]
-   --write-timeout value  The duration after which the client should timeout when writing to a connection (default: 5s) [$WRITE_TIMEOUT]
-   --log.format value     Specify the format of logs. Supported formats: 'logfmt', 'json', 'console' [$LOG_FORMAT]
-   --log.level value      Specify the log level. e.g. 'debug', 'info', 'error'. (default: "info") [$LOG_LEVEL]
-   --log.ctx value        A list of context field appended to every log. Format: key=value. [$LOG_CTX]
-   --help, -h             show help (default: false)
+   --protocol value                     Use protocol for the connection (available: tcp, udp) (default: "tcp") [$PROTOCOL]
+   --addr value                         The address of the connqc server [$ADDR]
+   --backoff value                      The duration to wait for before retrying to connect to the server (default: 1s) [$BACKOFF]
+   --interval value                     The interval at which to send probe messages to the server (default: 1s) [$INTERVAL]
+   --read-timeout value                 The duration after which the client should timeout when reading from a connection (default: 2s) [$READ_TIMEOUT]
+   --write-timeout value                The duration after which the client should timeout when writing to a connection (default: 5s) [$WRITE_TIMEOUT]
+   --log.format value                   Specify the format of logs. Supported formats: 'logfmt', 'json', 'console' [$LOG_FORMAT]
+   --log.level value                    Specify the log level. e.g. 'debug', 'info', 'error'. (default: "info") [$LOG_LEVEL]
+   --log.ctx value [ --log.ctx value ]  A list of context field appended to every log. Format: key=value. [$LOG_CTX]
+   --help, -h                           show help
+
 ```
 
 ## License

client.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/hamba/logger/v2"
 	lctx "github.com/hamba/logger/v2/ctx"
+	"github.com/nitrado/connqc/tcp"
+	"github.com/nitrado/connqc/udp"
 )
 
 // Client attempts to hold a connection with a server, sending probe messages at a configured interval.
@@ -33,11 +35,27 @@ func NewClient(backoff, sendInterval, readTimeout, writeTimeout time.Duration, l
 
 // Run sends probe messages to the server continuously.
 // If the connection fails, it retries at the configured backoff interval.
-func (c *Client) Run(ctx context.Context, addr string) {
+func (c *Client) Run(ctx context.Context, protocol, addr string) {
+	var (
+		conn net.Conn
+		err  error
+		idx  int
+	)
 	for {
-		conn, err := net.Dial("tcp", addr)
+		log := c.log.With(lctx.Str("protocol", protocol), lctx.Str("addr", addr), lctx.Int("reconnect", idx))
+		idx++
+
+		switch protocol {
+		case "tcp":
+			conn, err = tcp.Connect(addr)
+		case "udp":
+			conn, err = udp.Connect(addr)
+		default:
+			log.Error("Unexpected protocol")
+			return
+		}
 		if err != nil {
-			c.log.Error("Could not connect to server", lctx.Str("address", addr))
+			log.Error("Could not connect to server", lctx.Err(err))
 
 			select {
 			case <-ctx.Done():
@@ -48,7 +66,7 @@ func (c *Client) Run(ctx context.Context, addr string) {
 		}
 
 		if err = c.handleConn(ctx, conn); err != nil {
-			c.log.Error("Connection error", lctx.Error("error", err))
+			log.Error("Connection error", lctx.Err(err))
 		}
 
 		select {
@@ -64,7 +82,7 @@ type expectation struct {
 	probe     Probe
 }
 
-func (c *Client) handleConn(ctx context.Context, conn net.Conn) error { //nolint:funlen
+func (c *Client) handleConn(ctx context.Context, conn net.Conn) error { //nolint:funlen // Simplify readability.
 	defer func() { _ = conn.Close() }()
 
 	readCh := make(chan readResponse)
@@ -86,10 +104,10 @@ func (c *Client) handleConn(ctx context.Context, conn net.Conn) error { //nolint
 				Data: fmt.Sprintf("Hello %d", id),
 			}
 			if err := enc.Encode(p); err != nil {
-				return fmt.Errorf("writing Message: %w", err)
+				return fmt.Errorf("writing message: %w", err)
 			}
 
-			c.log.Debug("Sent probe", lctx.Interface("probe", p))
+			c.log.Info("Message sent", lctx.Interface("probe", p))
 
 			id++
 			expect = append(expect, expectation{timestamp: time.Now(), probe: p})
@@ -98,7 +116,7 @@ func (c *Client) handleConn(ctx context.Context, conn net.Conn) error { //nolint
 				return nil
 			}
 			if resp.err != nil {
-				return fmt.Errorf("reading Message: %w", resp.err)
+				return fmt.Errorf("reading response: %w", resp.err)
 			}
 
 			var (
@@ -115,7 +133,12 @@ func (c *Client) handleConn(ctx context.Context, conn net.Conn) error { //nolint
 					break
 				}
 
-				c.log.Warn("Message lost", lctx.Uint64("id", exp.probe.ID), lctx.Str("data", exp.probe.Data))
+				c.log.Warn("Message dropped",
+					lctx.Str("error", "unexpected ID"),
+					lctx.Uint64("expected_id", resp.probe.ID),
+					lctx.Uint64("id", exp.probe.ID),
+					lctx.Str("data", exp.probe.Data),
+				)
 			}
 			if !found {
 				c.log.Error("No expectation found")
@@ -152,12 +175,10 @@ func (c *Client) readLoop(conn net.Conn, ch chan readResponse) {
 
 		p, ok := msg.(Probe)
 		if !ok {
-			ch <- readResponse{err: err}
+			ch <- readResponse{err: fmt.Errorf("message not a probe: %T", msg)}
 			continue
 		}
 
-		c.log.Debug("Received probe", lctx.Interface("probe", p))
-
 		ch <- readResponse{timestamp: time.Now(), probe: p}
 	}
 }

cmd/connqc/client.go

@@ -1,28 +1,41 @@
 package main
 
 import (
+	"context"
+	"fmt"
+
 	"github.com/hamba/cmd/v2"
+	lctx "github.com/hamba/logger/v2/ctx"
 	"github.com/nitrado/connqc"
 	"github.com/urfave/cli/v2"
 )
 
 func runClient(c *cli.Context) error {
-	ctx := c.Context
+	ctx, cancel := context.WithCancel(c.Context)
+	defer cancel()
 
 	log, err := cmd.NewLogger(c)
 	if err != nil {
 		return err
 	}
 
+	protocol := c.String(flagProtocol)
+	if protocol != flagProtocolTCP && protocol != flagProtocolUDP {
+		return fmt.Errorf("unsupported protocol: %s", protocol)
+	}
+
 	backoff := c.Duration(flagConnBackoff)
 	sendInterval := c.Duration(flagSendInterval)
 	readTimeout := c.Duration(flagReadTimeout)
 	writeTimeout := c.Duration(flagWriteTimeout)
 
-	sig := connqc.NewClient(backoff, sendInterval, readTimeout, writeTimeout, log)
+	log = log.With(lctx.Str("protocol", protocol))
 
-	addr := c.String(flagAddr)
-	go sig.Run(ctx, addr)
+	client := connqc.NewClient(backoff, sendInterval, readTimeout, writeTimeout, log)
+	go func() {
+		client.Run(ctx, protocol, c.String(flagAddr))
+		cancel()
+	}()
 
 	<-ctx.Done()
 

cmd/connqc/main.go

@@ -16,7 +16,10 @@ import (
 )
 
 const (
-	flagAddr = "addr"
+	flagProtocol    = "protocol"
+	flagProtocolTCP = "tcp"
+	flagProtocolUDP = "udp"
+	flagAddr        = "addr"
 
 	flagBufferSize   = "buffer-size"
 	flagReadTimeout  = "read-timeout"
@@ -33,21 +36,29 @@ var commands = []*cli.Command{
 		Name:  "client",
 		Usage: "Run the connqc client",
 		Flags: cmd.Flags{
+			&cli.StringFlag{
+				Name: flagProtocol,
+				Usage: fmt.Sprintf(
+					"The protocol for the connection. Supported protocols: '%s', '%s'", flagProtocolTCP, flagProtocolUDP,
+				),
+				Value:   flagProtocolTCP,
+				EnvVars: []string{strcase.ToSNAKE(flagProtocol)},
+			},
 			&cli.StringFlag{
 				Name:     flagAddr,
-				Usage:    "The address of a connqc server",
+				Usage:    "The address of the connqc server",
 				Required: true,
 				EnvVars:  []string{strcase.ToSNAKE(flagAddr)},
 			},
 			&cli.DurationFlag{
 				Name:    flagConnBackoff,
-				Usage:   "Duration to wait for before retrying to connect to the server",
+				Usage:   "The duration to wait for before retrying to connect to the server",
 				Value:   time.Second,
 				EnvVars: []string{strcase.ToSNAKE(flagConnBackoff)},
 			},
 			&cli.DurationFlag{
 				Name:    flagSendInterval,
-				Usage:   "Interval at which to send probe messages to the server",
+				Usage:   "The interval at which to send probe messages to the server",
 				Value:   time.Second,
 				EnvVars: []string{strcase.ToSNAKE(flagSendInterval)},
 			},

cmd/connqc/server.go

@@ -2,15 +2,18 @@ package main
 
 import (
 	"errors"
+	"net"
+	"sync"
 
 	"github.com/hamba/cmd/v2"
 	lctx "github.com/hamba/logger/v2/ctx"
 	"github.com/nitrado/connqc"
 	"github.com/nitrado/connqc/tcp"
+	"github.com/nitrado/connqc/udp"
 	"github.com/urfave/cli/v2"
 )
 
-func runServer(c *cli.Context) error {
+func runServer(c *cli.Context) error { //nolint:funlen // Keep it simple and readable.
 	ctx := c.Context
 
 	log, err := cmd.NewLogger(c)
@@ -22,34 +25,56 @@ func runServer(c *cli.Context) error {
 	readTimeout := c.Duration(flagReadTimeout)
 	writeTimeout := c.Duration(flagWriteTimeout)
 
-	sigSrv := connqc.NewServer(bufferSize, readTimeout, writeTimeout, log)
+	srv := connqc.NewServer(bufferSize, readTimeout, writeTimeout, log)
 
-	srv, err := tcp.NewServer(sigSrv)
+	tcpSrv, err := tcp.NewServer(srv)
+	if err != nil {
+		return err
+	}
+
+	udpSrv, err := udp.NewServer(srv)
 	if err != nil {
 		return err
 	}
 
 	addr := c.String(flagAddr)
+
+	grp := sync.WaitGroup{}
+
 	log.Info("Starting server",
 		lctx.Str("addr", addr),
 		lctx.Int("buffer_size", bufferSize),
 		lctx.Duration("read_timeout", readTimeout),
 		lctx.Duration("write_timeout", writeTimeout),
 	)
+
+	grp.Add(1)
 	go func() {
-		if err := srv.Listen(addr); err != nil && !errors.Is(err, tcp.ErrServerClosed) {
-			log.Error("Server error", lctx.Error("error", err))
+		defer grp.Done()
+
+		if err := tcpSrv.Listen(ctx, addr); err != nil && !errors.Is(err, net.ErrClosed) {
+			log.Error("Server error", lctx.Str("protocol", "tcp"), lctx.Err(err))
+			return
 		}
+		log.Info("TCP server stopped")
+	}()
+
+	grp.Add(1)
+	go func() {
+		defer grp.Done()
+
+		if err := udpSrv.Listen(ctx, addr); err != nil && !errors.Is(err, net.ErrClosed) {
+			log.Error("Server error", lctx.Str("protocol", "udp"), lctx.Err(err))
+			return
+		}
+		log.Info("UDP server stopped")
 	}()
-	defer func() { _ = srv.Close() }()
 
 	<-ctx.Done()
 
 	log.Info("Shutting down")
 
-	if err = srv.Close(); err != nil {
-		log.Warn("Failed to shutdown server", lctx.Error("error", err))
-	}
+	grp.Wait()
 
 	return nil
 }