neoirc/internal/ircserver/server.go

package ircserver

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"go.uber.org/fx"
	"sneak.berlin/go/neoirc/internal/broker"
	"sneak.berlin/go/neoirc/internal/config"
	"sneak.berlin/go/neoirc/internal/db"
	"sneak.berlin/go/neoirc/internal/logger"
	"sneak.berlin/go/neoirc/internal/service"
)

// stopTimeout is the longest stop waits for the connections'
// cleanup, which ends their sessions.
const stopTimeout = 5 * time.Second

// Params defines the dependencies for creating an IRC
// Server.
type Params struct {
	fx.In

	Logger   *logger.Logger
	Config   *config.Config
	Database *db.Database
	Broker   *broker.Broker
	Service  *service.Service
}

// Server is the TCP IRC protocol server.
type Server struct {
	log      *slog.Logger
	cfg      *config.Config
	database *db.Database
	brk      *broker.Broker
	svc      *service.Service
	listener net.Listener
	mu       sync.Mutex
	conns    map[*Conn]struct{}
	cancel   context.CancelFunc

	// serving counts the connections whose serve, cleanup
	// included, has not returned yet.
	serving sync.WaitGroup

	// pingInterval and pongDeadline are the constants of the
	// same names, except in tests, which shorten them.
	pingInterval time.Duration
	pongDeadline time.Duration

	// pollInterval and stopTimeout are the constants of the
	// same names, except in tests, which change them.
	pollInterval time.Duration
	stopTimeout  time.Duration

	// resolveHost is the function of the same name, except in
	// tests, which replace it with one that blocks.
	resolveHost func(ctx context.Context, addr string) string
}

// New creates a new IRC Server and registers its lifecycle
// hooks.  The listener is only started if IRC_LISTEN_ADDR
// is configured; otherwise the server is inert.
func New(
	lifecycle fx.Lifecycle,
	params Params,
) *Server {
	srv := &Server{
		log:      params.Logger.Get(),
		cfg:      params.Config,
		database: params.Database,
		brk:      params.Broker,
		svc:      params.Service,
		conns:    make(map[*Conn]struct{}),
		listener: nil,
		cancel:   nil,
		mu:       sync.Mutex{},
		serving:  sync.WaitGroup{},

		pingInterval: pingInterval,
		pongDeadline: pongDeadline,
		pollInterval: pollInterval,
		stopTimeout:  stopTimeout,
		resolveHost:  resolveHost,
	}

	listenAddr := params.Config.IRCListenAddr
	if listenAddr == "" {
		return srv
	}

	lifecycle.Append(fx.Hook{
		OnStart: func(ctx context.Context) error {
			return srv.start(ctx, listenAddr)
		},
		OnStop: func(_ context.Context) error {
			srv.stop()

			return nil
		},
	})

	return srv
}

// start begins listening for TCP connections.
//
//nolint:contextcheck // long-lived server ctx, not the short Fx one
func (s *Server) start(ctx context.Context, addr string) error {
	var listenConfig net.ListenConfig

	ln, err := listenConfig.Listen(ctx, "tcp", addr)
	if err != nil {
		return fmt.Errorf("irc listen: %w", err)
	}

	s.listener = ln

	serverCtx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	s.log.Info(
		"irc server listening", "addr", addr,
	)

	go s.acceptLoop(serverCtx)

	return nil
}

// stop shuts down the listener and all connections, then
// waits up to stopTimeout for every connection's cleanup to
// finish, so that their sessions are ended before the
// database closes.
func (s *Server) stop() {
	if s.cancel != nil {
		s.cancel()
	}

	if s.listener != nil {
		s.listener.Close() //nolint:errcheck,gosec
	}

	s.mu.Lock()
	for c := range s.conns {
		c.conn.Close() //nolint:errcheck,gosec
	}
	s.mu.Unlock()

	done := make(chan struct{})

	go func() {
		s.serving.Wait()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(s.stopTimeout):
		s.log.Error(
			"irc connections not cleaned up before stop timeout",
			"timeout", s.stopTimeout,
		)
	}
}

// acceptLoop accepts new connections.
func (s *Server) acceptLoop(ctx context.Context) {
	for {
		tcpConn, err := s.listener.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				s.log.Error(
					"irc accept error", "error", err,
				)

				continue
			}
		}

		client := newConn(
			tcpConn, s.log,
			s.database, s.brk, s.cfg, s.svc, s.pollInterval,
		)

		s.mu.Lock()

		// stop cancels ctx before it takes s.mu to close the
		// connections, so a connection accepted while stop
		// runs is either closed and waited for by stop or
		// closed here.
		if ctx.Err() != nil {
			s.mu.Unlock()
			tcpConn.Close() //nolint:errcheck,gosec

			return
		}

		s.conns[client] = struct{}{}
		s.serving.Add(1)
		s.mu.Unlock()

		go func() {
			defer s.serving.Done()
			defer func() {
				s.mu.Lock()
				delete(s.conns, client)
				s.mu.Unlock()
			}()

			client.serve(
				ctx, s.resolveHost,
				s.pingInterval, s.pongDeadline,
			)
		}()
	}
}