Add TCP Check

README.md

@@ -15,6 +15,7 @@
  * Memcached
  * InfluxDB
  * Nats
+ * TCP
 
 ## Usage
 

_examples/server.go

@@ -10,6 +10,7 @@ import (
  healthHttp "github.com/hellofresh/health-go/v5/checks/http"
  healthMySql "github.com/hellofresh/health-go/v5/checks/mysql"
  healthPg "github.com/hellofresh/health-go/v5/checks/postgres"
+ healthTcp "github.com/hellofresh/health-go/v5/checks/tcp"
 )
 
 func main() {
@@ -38,6 +39,17 @@ func main() {
  }),
  })
 
+ // tcp health check example
+ h.Register(health.Config{
+ Name: "tcp-check",
+ Timeout: time.Second * 5,
+ SkipOnErr: true,
+ Check: healthTcp.New(healthTcp.Config{
+ Host: `example.com`,
+ Port: 1123,
+ }),
+ })
+
  // postgres health check example
  h.Register(health.Config{
  Name: "postgres-check",

checks/tcp/check.go

@@ -0,0 +1,43 @@
+package tcp
+
+import (
+ "context"
+ "fmt"
+ "net"
+ "strconv"
+ "time"
+)
+
+const defaultRequestTimeout = 5 * time.Second
+
+// Config is the TCP checker configuration settings container.
+type Config struct {
+ // Host is the remote service hostname health check TCP.
+ Host string
+ // Port is the remote service port health check TCP.
+ Port int
+ // RequestTimeout is the duration that health check will try to consume published test message.
+ // If not set - 5 seconds
+ RequestTimeout time.Duration
+}
+
+// New creates new HTTP service health check that verifies the following:
+// - connection establishing
+// - getting response status from defined URL
+// - verifying that status code is less than 500
+func New(config Config) func(ctx context.Context) error {
+ if config.RequestTimeout == 0 {
+ config.RequestTimeout = defaultRequestTimeout
+ }
+
+ return func(ctx context.Context) error {
+ conn, err := net.DialTimeout("tcp", config.Host+":"+strconv.Itoa(config.Port), config.RequestTimeout)
+ if err != nil {
+ return fmt.Errorf("making the request for the health check failed: %w", err.Error())
+ }
+
+ conn.Close()
+
+ return nil
+ }
+}

checks/tcp/check_test.go

@@ -0,0 +1,53 @@
+package tcp
+
+import (
+ "github.com/stretchr/testify/require"
+ "log"
+ "net"
+ "strconv"
+ "testing"
+)
+
+var tcp Server
+
+func init() {
+ // Start the new server
+ tcp, err := NewServer("tcp", ":1123")
+ if err != nil {
+ log.Println("error starting TCP server")
+ return
+ }
+
+ // Run the servers in goroutines to stop blocking
+ go func() {
+ tcp.Run()
+ }()
+}
+
+func TestNew(t *testing.T) {
+
+ t.Run("service is available", func(t *testing.T) {
+
+ check := Config{
+ Host: "127.0.0.1",
+ Port: 1123,
+ RequestTimeout: defaultRequestTimeout,
+ }
+ conn, err := net.DialTimeout("tcp", check.Host+":"+strconv.Itoa(check.Port), check.RequestTimeout)
+ require.NoError(t, err)
+ defer conn.Close()
+ })
+
+ t.Run("service is not available", func(t *testing.T) {
+ check := Config{
+ Host: "127.0.0.1",
+ Port: 1124,
+ RequestTimeout: defaultRequestTimeout,
+ }
+
+ conn, err := net.DialTimeout("tcp", check.Host+":"+strconv.Itoa(check.Port), check.RequestTimeout)
+ require.Error(t, err)
+ defer conn.Close()
+
+ })
+}

checks/tcp/net.go

@@ -0,0 +1,152 @@
+package tcp
+
+import (
+ "bufio"
+ "errors"
+ "fmt"
+ "log"
+ "net"
+ "strings"
+)
+
+// Server defines the minimum contract our
+// TCP and UDP server implementations must satisfy.
+type Server interface {
+ Run() error
+ Close() error
+}
+
+// NewServer creates a new Server using given protocol
+// and addr.
+func NewServer(protocol, addr string) (Server, error) {
+ switch strings.ToLower(protocol) {
+ case "tcp":
+ return &TCPServer{
+ addr: addr,
+ }, nil
+ case "udp":
+ return &UDPServer{
+ addr: addr,
+ }, nil
+ }
+ return nil, errors.New("Invalid protocol given")
+}
+
+// TCPServer holds the structure of our TCP
+// implementation.
+type TCPServer struct {
+ addr string
+ server net.Listener
+}
+
+// Run starts the TCP Server.
+func (t *TCPServer) Run() (err error) {
+ t.server, err = net.Listen("tcp", t.addr)
+ if err != nil {
+ return err
+ }
+ defer t.Close()
+
+ for {
+ conn, err := t.server.Accept()
+ if err != nil {
+ err = errors.New("could not accept connection")
+ break
+ }
+ if conn == nil {
+ err = errors.New("could not create connection")
+ break
+ }
+ return t.handleConnections()
+ }
+ return
+}
+
+// Close shuts down the TCP Server
+func (t *TCPServer) Close() (err error) {
+ return t.server.Close()
+}
+
+// handleConnections is used to accept connections on
+// the TCPServer and handle each of them in separate
+// goroutines.
+func (t *TCPServer) handleConnections() (err error) {
+ for {
+ conn, err := t.server.Accept()
+ if err != nil || conn == nil {
+ err = errors.New("could not accept connection")
+ break
+ }
+
+ go t.handleConnection(conn)
+ }
+ return
+}
+
+// handleConnections deals with the business logic of
+// each connection and their requests.
+func (t *TCPServer) handleConnection(conn net.Conn) {
+ defer conn.Close()
+
+ rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
+ for {
+ req, err := rw.ReadString('\n')
+ if err != nil {
+ rw.WriteString("failed to read input")
+ rw.Flush()
+ return
+ }
+
+ rw.WriteString(fmt.Sprintf("Request received: %s", req))
+ rw.Flush()
+ }
+}
+
+// UDPServer holds the necessary structure for our
+// UDP server.
+type UDPServer struct {
+ addr string
+ server *net.UDPConn
+}
+
+// Run starts the UDP server.
+func (u *UDPServer) Run() (err error) {
+ laddr, err := net.ResolveUDPAddr("udp", u.addr)
+ if err != nil {
+ return errors.New("could not resolve UDP addr")
+ }
+
+ u.server, err = net.ListenUDP("udp", laddr)
+ if err != nil {
+ return errors.New("could not listen on UDP")
+ }
+
+ return u.handleConnections()
+}
+
+func (u *UDPServer) handleConnections() error {
+ var err error
+ for {
+ buf := make([]byte, 2048)
+ n, conn, err := u.server.ReadFromUDP(buf)
+ if err != nil {
+ log.Println(err)
+ break
+ }
+ if conn == nil {
+ continue
+ }
+
+ go u.handleConnection(conn, buf[:n])
+ }
+ return err
+}
+
+func (u *UDPServer) handleConnection(addr *net.UDPAddr, cmd []byte) {
+ u.server.WriteToUDP([]byte(fmt.Sprintf("Request recieved: %s", cmd)), addr)
+}
+
+// Close ensures that the UDPServer is shut down gracefully.
+func (u *UDPServer) Close() error {
+ return u.server.Close()
+}

docs/index.md

@@ -13,6 +13,7 @@
  * gRPC
  * Memcached
  * Nats
+ * TCP
 
 ## Usage