netris-nestri/packages/cli/internal/party/client.go

package party

import (
	"fmt"
	"nestrilabs/cli/internal/machine"
	"nestrilabs/cli/internal/resource"
	"net/http"
	"net/url"
	"time"

	"github.com/charmbracelet/log"
	"github.com/gorilla/websocket"
)

const (
	// Initial retry delay
	initialRetryDelay = 1 * time.Second
	// Maximum retry delay
	maxRetryDelay = 30 * time.Second
	// Factor to increase delay by after each attempt
	backoffFactor = 2
)

type Party struct {
	// Channel to signal shutdown
	done        chan struct{}
	fingerprint string
	hostname    string
}

func NewParty() *Party {
	m := machine.NewMachine()
	fingerpint := m.GetMachineID()
	return &Party{
		done:        make(chan struct{}),
		fingerprint: fingerpint,
		hostname:    m.Hostname,
	}
}

// Shutdown gracefully closes the connection
func (p *Party) Shutdown() {
	close(p.done)
}

func (p *Party) Connect() {
	baseURL := fmt.Sprintf("ws://localhost:1999/parties/main/%s", p.fingerprint)
	params := url.Values{}
	params.Add("_pk", p.hostname)
	wsURL := baseURL + "?" + params.Encode()

	retryDelay := initialRetryDelay
	header := http.Header{}
	bearer := fmt.Sprintf("Bearer %s", resource.Resource.AuthFingerprintKey.Value)
	header.Add("Authorization", bearer)

	for {
		select {
		case <-p.done:
			log.Info("Shutting down connection")
			return
		default:
			conn, _, err := websocket.DefaultDialer.Dial(wsURL, header)
			if err != nil {
				log.Error("Failed to connect to party server", "err", err)
				time.Sleep(retryDelay)
				// Increase retry delay exponentially, but cap it
				retryDelay = time.Duration(float64(retryDelay) * backoffFactor)
				if retryDelay > maxRetryDelay {
					retryDelay = maxRetryDelay
				}
				continue
			}
			log.Info("Connection to server", "url", wsURL)

			// Reset retry delay on successful connection
			retryDelay = initialRetryDelay

			// Handle connection in a separate goroutine
			connectionClosed := make(chan struct{})
			go func() {
				defer close(connectionClosed)
				defer conn.Close()

				// Send initial message
				// if err := conn.WriteMessage(websocket.TextMessage, []byte("hello there")); err != nil {
				// 	log.Error("Failed to send initial message", "err", err)
				// 	return
				// }

				// Read messages loop
				for {
					select {
					case <-p.done:
						return
					default:
						_, message, err := conn.ReadMessage()
						if err != nil {
							log.Error("Error reading message", "err", err)
							return
						}
						log.Info("Received message from party server", "message", string(message))
					}
				}
			}()

			// Wait for either connection to close or shutdown signal
			select {
			case <-connectionClosed:
				log.Warn("Connection closed, attempting to reconnect...")
				time.Sleep(retryDelay)
			case <-p.done:
				log.Info("Shutting down connection")
				return
			}
		}
	}
}