HomesteadGateway/ws/registry.go

package ws

import (
	"fmt"
	"homestead/homestead_gateway/util"
	"time"

	"github.com/gorilla/websocket"
)

func (q *BoundedQueue) Enqueue(m GatewayMessageOut) bool {
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.capacity == 0 {
		return false
	}

	if q.length < q.capacity {
		q.buf[(q.start+q.length)%q.capacity] = m
		q.length++
		return true
	}

	// overwrite oldest
	q.buf[q.start] = m
	q.start = (q.start + 1) % q.capacity

	return true
}

func (q *BoundedQueue) PopAll() []GatewayMessageOut {
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.length == 0 {
		return nil
	}

	out := make([]GatewayMessageOut, 0, q.length)
	for i := 0; i < q.length; i++ {
		out = append(out, q.buf[(q.start+i)%q.capacity])
	}

	q.start = 0
	q.length = 0

	return out
}

func (q *BoundedQueue) Len() int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.length
}

//

func (r *Registry) getOrCreate(channel string) *ChannelEntry {
	r.mu.RLock()
	e := r.entries[channel]
	r.mu.RUnlock()
	if e != nil {
		return e
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if e = r.entries[channel]; e == nil {
		e = newChannelEntry(channel, r.queueCap)
		r.entries[channel] = e
	}
	return e
}

func (r *Registry) ForEach(cb func(channelID string)) {
	r.mu.RLock()
	ids := make([]string, 0, len(r.entries))
	for id := range r.entries {
		ids = append(ids, id)
	}
	r.mu.RUnlock()

	for _, id := range ids {
		cb(id)
	}
}

//

// RegisterMod : map channel_id -> mod conn (serverID)
func (r *Registry) RegisterMod(channelID, serverID string, conn *websocket.Conn) {
	e := r.getOrCreate(channelID)
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.Mod != nil && e.Mod.Conn != nil {
		util.CloseConn(e.Mod.Conn)
	}

	e.Mod = &ConnWrapper{Conn: conn, ServerID: serverID, LastSeen: time.Now()}
	// flush queued bot->mod messages for this channel
	// caller should use FlushChannelWithSender to perform actual sends
}

// RegisterBot : single connection for bot. after registration call FlushAllToBotWithSender
func (r *Registry) RegisterBot(conn *websocket.Conn) {
	r.botMu.Lock()

	if r.bot != nil && r.bot.Conn != nil {
		r.UnregisterBot()
	}

	r.bot = &ConnWrapper{Conn: conn, LastSeen: time.Now()}
	r.botMu.Unlock()
}

func (r *Registry) UnregisterMod(channelID string) {
	r.mu.RLock()
	e := r.entries[channelID]
	r.mu.RUnlock()
	if e == nil {
		return
	}

	e.mu.Lock()
	modConn := e.Mod
	e.Mod = nil
	e.mu.Unlock()

	if modConn != nil && modConn.Conn != nil {
		util.CloseConn(modConn.Conn)
	}
}

func (r *Registry) UnregisterBot() {
	r.botMu.Lock()
	botConn := r.bot
	r.bot = nil
	r.botMu.Unlock()

	if botConn != nil && botConn.Conn != nil {
		util.CloseConn(botConn.Conn)
	}
}

func (r *Registry) Send(channelID string, out GatewayMessageOut, sendOverConn func(*websocket.Conn, GatewayMessageOut) error) (delivered bool, queued bool, err error) {
	if out.Type == "mod" {
		r.botMu.Lock()
		b := r.bot
		r.botMu.Unlock()

		if b != nil && b.Conn != nil {
			if err := sendOverConn(b.Conn, out); err == nil {
				return true, false, nil
			}
			r.UnregisterBot()
		}

		e := r.getOrCreate(channelID)
		e.mu.Lock()
		enq := e.Queue.Enqueue(out)
		e.mu.Unlock()

		if !enq {
			return false, false, fmt.Errorf("queue disabled")
		}
		return false, true, nil
	}

	e := r.getOrCreate(channelID)
	e.mu.Lock()
	mod := e.Mod
	e.mu.Unlock()

	if mod != nil && mod.Conn != nil {
		if err := sendOverConn(mod.Conn, out); err == nil {
			return true, false, nil
		}
		r.UnregisterMod(channelID)
	}

	e.mu.Lock()
	enq := e.Queue.Enqueue(out)
	e.mu.Unlock()

	if !enq {
		return false, false, fmt.Errorf("queue disabled")
	}
	return false, true, nil
}

//

func (r *Registry) FlushChannelWithSender(channelID string, sendOverConn func(*websocket.Conn, GatewayMessageOut) error) {
	r.mu.RLock()
	e := r.entries[channelID]
	r.mu.RUnlock()
	if e == nil {
		return
	}

	e.mu.Lock()
	if e.Mod == nil || e.Mod.Conn == nil {
		e.mu.Unlock()
		return
	}

	msgs := e.Queue.PopAll()
	modConn := e.Mod.Conn
	e.mu.Unlock()

	for _, m := range msgs {
		if err := sendOverConn(modConn, m); err != nil {
			// if send fails, re-enqueue (best-effort), drop-oldest logic applies
			e.mu.Lock()
			_ = e.Queue.Enqueue(m)
			e.mu.Unlock()
		}
	}
}

func (r *Registry) FlushAllToBotWithSender(sendOverConn func(*websocket.Conn, GatewayMessageOut) error) {
	r.botMu.Lock()
	b := r.bot
	r.botMu.Unlock()
	if b == nil || b.Conn == nil {
		return
	}

	r.mu.RLock()
	entries := make([]*ChannelEntry, 0, len(r.entries))
	for _, e := range r.entries {
		entries = append(entries, e)
	}
	r.mu.RUnlock()

	for _, e := range entries {
		e.mu.Lock()
		msgs := e.Queue.PopAll()
		e.mu.Unlock()

		if len(msgs) == 0 {
			continue
		}
		for _, m := range msgs {
			if err := sendOverConn(b.Conn, m); err != nil {
				e.mu.Lock()
				_ = e.Queue.Enqueue(m)
				e.mu.Unlock()
			}
		}
	}
}