TCP Socket Example in Golang

main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"time"

	"golang.org/x/sync/errgroup"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()

	eg, ctx := errgroup.WithContext(ctx)

	startedCh := make(chan struct{})
	port := 8082

	eg.Go(func() error {
		return startServer(ctx, port, startedCh)
	})

	eg.Go(func() error {
		return startClient(port, startedCh)
	})

	if err := eg.Wait(); err != nil {
		return fmt.Errorf("failed to wait for goroutines: %w", err)
	}

	return nil
}

func startServer(ctx context.Context, port int, startedCh chan struct{}) error {
	lc := net.ListenConfig{}

	ln, err := lc.Listen(ctx, "tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return fmt.Errorf("failed to listen: %w", err)
	}

	log.Printf("listening on %s", ln.Addr())

	close(startedCh)

	var conn net.Conn

	conn, err = ln.Accept()
	if err != nil {
		return fmt.Errorf("failed to accept connection: %w", err)
	}

	log.Printf("accepted connection from %s", conn.RemoteAddr())

	if err = handleConnection(ctx, conn); err != nil {
		return fmt.Errorf("failed to handle connection: %w", err)
	}

	log.Printf("handled connection from %s", conn.RemoteAddr())

	return nil
}

func startClient(port int, startedCh <-chan struct{}) error {
	<-startedCh

	address := fmt.Sprintf("127.0.0.1:%d", port)

	log.Printf("dialing %s", address)

	conn, err := net.DialTimeout("tcp", address, 3*time.Second)
	if err != nil {
		return fmt.Errorf("failed to dial: %w", err)
	}

	log.Printf("dialed %s", address)

	readTimeout := 10 * time.Second
	if err = conn.SetDeadline(time.Now().Add(readTimeout)); err != nil {
		return fmt.Errorf("failed to set read deadline: %w", err)
	}

	log.Printf("set initial read timeout to %s", readTimeout)

	numRead := 0
	buffer := make([]byte, 1024)

	for {
		numRead, err = conn.Read(buffer)
		if err != nil {
			if err == io.EOF {
				log.Printf("connection closed by the remote")

				return nil
			}

			return fmt.Errorf("failed to read from connection: %w", err)
		}

		log.Printf("read %d bytes: %s", numRead, string(buffer[:numRead]))

		// set a new deadline after every read.
		// if you comment this out, the connection will be closed after the read timeout.
		if numRead > 0 {
			if err = conn.SetDeadline(time.Now().Add(readTimeout)); err != nil {
				return fmt.Errorf("failed to set read deadline after read: %w", err)
			}
		}
	}
}

func handleConnection(ctx context.Context, conn net.Conn) (err error) {
	defer func() {
		if closeErr := conn.Close(); closeErr != nil {
			err = errors.Join(err, fmt.Errorf("failed to close connection: %w", closeErr))
		}
	}()

	ticker := time.NewTicker(4 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			if errors.Is(ctx.Err(), context.Canceled) {
				return nil
			}

			return fmt.Errorf("context error: %w", ctx.Err())
		case <-ticker.C:
			// write a line to the connection
			data := []byte(fmt.Sprintf("hello %d\n", time.Now().UnixMilli()))

			if _, err = conn.Write(data); err != nil {
				return fmt.Errorf("failed to write to connection: %w", err)
			}

			log.Printf("sent len(%d): %s", len(data), string(data))
		}
	}
}