go-honeybee/transport/connection.go

package transport

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

	"git.wisehodl.dev/jay/go-honeybee/types"
	"github.com/gorilla/websocket"
)

type ConnectionState int

const (
	StateDisconnected ConnectionState = iota
	StateConnecting
	StateConnected
	StateClosed
)

func (s ConnectionState) String() string {
	switch s {
	case StateDisconnected:
		return "disconnected"
	case StateConnecting:
		return "connecting"
	case StateConnected:
		return "connected"
	case StateClosed:
		return "closed"
	default:
		return "unknown"
	}
}

type Connection struct {
	url    *url.URL
	dialer types.Dialer
	socket types.Socket
	config *ConnectionConfig
	logger *slog.Logger

	incoming chan []byte
	errors   chan error
	done     chan struct{}

	state ConnectionState

	wg          sync.WaitGroup
	closed      bool
	mu          sync.RWMutex
	writeMu     sync.Mutex
	doneOnce    sync.Once
	cleanupOnce sync.Once
}

func NewConnection(urlStr string, config *ConnectionConfig, logger *slog.Logger) (*Connection, error) {
	if config == nil {
		config = GetDefaultConnectionConfig()
	}

	if err := ValidateConnectionConfig(config); err != nil {
		return nil, err
	}

	url, err := ParseURL(urlStr)
	if err != nil {
		return nil, err
	}

	conn := &Connection{
		url:      url,
		dialer:   NewDialer(),
		socket:   nil,
		config:   config,
		logger:   logger,
		incoming: make(chan []byte, 100),
		errors:   make(chan error, 10),
		state:    StateDisconnected,
		done:     make(chan struct{}),
	}

	return conn, nil
}

func NewConnectionFromSocket(
	socket types.Socket, config *ConnectionConfig, logger *slog.Logger,
) (*Connection, error) {
	if socket == nil {
		return nil, NewConnectionError(ErrNilSocket)
	}

	if config == nil {
		config = GetDefaultConnectionConfig()
	}

	if err := ValidateConnectionConfig(config); err != nil {
		return nil, err
	}

	conn := &Connection{
		url:      nil,
		dialer:   nil,
		socket:   socket,
		config:   config,
		logger:   logger,
		incoming: make(chan []byte, 100),
		errors:   make(chan error, 10),
		state:    StateConnected,
		done:     make(chan struct{}),
	}

	if config.CloseHandler != nil {
		socket.SetCloseHandler(config.CloseHandler)
	}

	conn.startReader()

	return conn, nil
}

func (c *Connection) Connect(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.socket != nil {
		return NewConnectionError(ErrSocketExists)
	}

	if c.closed {
		return NewConnectionError(ErrConnectionClosed)
	}

	if c.logger != nil {
		c.logger.Info("connecting")
	}

	c.state = StateConnecting

	retryMgr := NewRetryManager(c.config.Retry)
	socket, _, err := AcquireSocket(
		ctx, retryMgr, c.dialer, c.url.String(), c.logger)

	if err != nil {
		c.state = StateDisconnected
		if c.logger != nil {
			c.logger.Error("connection failed", "error", err)
		}
		return NewConnectionError(err)
	}

	c.socket = socket
	c.state = StateConnected

	if c.config.CloseHandler != nil {
		c.socket.SetCloseHandler(c.config.CloseHandler)
	}

	if c.logger != nil {
		c.logger.Info("connected")
	}

	c.startReader()

	return nil
}

func (c *Connection) Close() {
	c.shutdownExternal()
}

func (c *Connection) shutdownExternal() {
	err := c.shutdownSetClosed(true)
	if err != nil {
		return
	}

	c.shutdownInner()
	c.shutdownCleanup()
}

func (c *Connection) shutdownInternal() {
	err := c.shutdownSetClosed(false)
	if err != nil {
		return
	}

	c.shutdownInner()

	// defer final cleanup to allow this function to return
	// otherwise, a deadlock occurs where startReader triggers a shutdown and
	// must wait for itself to exit.
	go func() {
		c.shutdownCleanup()
	}()
}

func (c *Connection) shutdownInner() {
	c.shutdownSignalDone()
	c.shutdownLogStart()
	c.shutdownCloseSocket()
}

func (c *Connection) shutdownCleanup() {
	c.cleanupOnce.Do(func() {
		c.wg.Wait()
		c.shutdownCloseChannels()
		c.shutdownLogComplete()
	})
}

func (c *Connection) shutdownSetClosed(wait bool) error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return NewConnectionError(ErrConnectionClosed)
	}
	c.closed = true
	c.state = StateClosed
	c.mu.Unlock()
	return nil
}

func (c *Connection) shutdownSignalDone() {
	c.doneOnce.Do(func() {
		close(c.done)
	})
}

func (c *Connection) shutdownLogStart() {
	if c.logger != nil {
		c.logger.Info("shutting down")
	}
}

func (c *Connection) shutdownCloseSocket() {
	if c.socket != nil {
		// force unblock of any network operations immediately
		expired := time.Now().Add(-1 * time.Minute)
		c.socket.SetReadDeadline(expired)
		c.socket.SetWriteDeadline(expired)

		// close socket
		err := c.socket.Close()

		if err != nil && c.logger != nil {
			c.logger.Error("socket close failed", "error", err)
		}
	}
}

func (c *Connection) shutdownCloseChannels() {
	close(c.incoming)
	close(c.errors)
}

func (c *Connection) shutdownLogComplete() {
	if c.logger != nil {
		c.logger.Info("connection closed")
	}
}

func (c *Connection) startReader() {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		defer c.shutdownInternal()

		for {
			select {
			case <-c.done:
				return
			default:
				messageType, data, err := c.socket.ReadMessage()
				if err != nil {
					var wrappedErr error
					var closeErr *websocket.CloseError
					if errors.As(err, &closeErr) {
						switch closeErr.Code {
						case websocket.CloseNormalClosure, websocket.CloseGoingAway:
							if c.logger != nil {
								c.logger.Info("connection closed by peer",
									"code", closeErr.Code,
									"text", closeErr.Text,
								)
							}
							wrappedErr = fmt.Errorf("%w: %w", ErrPeerClosedClean, err)
						default:
							if c.logger != nil {
								c.logger.Error("unexpected close",
									"code", closeErr.Code,
									"text", closeErr.Text,
								)
							}
							wrappedErr = fmt.Errorf("%w: %w", ErrPeerClosedUnexpected, err)
						}
					} else {
						if c.logger != nil {
							c.logger.Error("read error", "error", err)
						}
						wrappedErr = fmt.Errorf("%w: %w", ErrReadError, err)
					}

					select {
					case <-c.done:
					case c.errors <- wrappedErr:
					}
					return
				}

				if messageType == websocket.TextMessage ||
					messageType == websocket.BinaryMessage {
					select {
					case <-c.done:
						return
					case c.incoming <- data:
					}
				}

			}
		}
	}()
}

func (c *Connection) Send(data []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	if c.closed {
		return NewConnectionError(ErrConnectionClosed)
	}

	if c.config.WriteTimeout > 0 {
		if err := c.socket.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout)); err != nil {
			if c.logger != nil {
				c.logger.Error("write deadline error", "error", err)
			}
			c.shutdownExternal()
			return NewConnectionError(fmt.Errorf("%w: %w", ErrFailedWriteDeadline, err))
		}
	}

	if err := c.socket.WriteMessage(websocket.TextMessage, data); err != nil {
		if c.logger != nil {
			c.logger.Error("write error", "error", err)
		}
		return NewConnectionError(fmt.Errorf("%w: %w", ErrWriteFailed, err))
	}

	return nil
}

func (c *Connection) Incoming() <-chan []byte {
	return c.incoming
}

func (c *Connection) Errors() <-chan error {
	return c.errors
}

func (c *Connection) State() ConnectionState {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.state
}

func (c *Connection) SetDialer(d types.Dialer) {
	c.dialer = d
}