aiforge/services/socketwrap/clientManager.go

package socketwrap

import (
	"os"
	"os/signal"
	"syscall"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
)

type ClientsManager struct {
	Clients    map[*Client]bool
	Register   chan *Client
	Unregister chan *Client
}

func NewClientsManager() *ClientsManager {
	return &ClientsManager{
		Register:   make(chan *Client),
		Unregister: make(chan *Client),
		Clients:    make(map[*Client]bool),
	}
}

var LastActionsQueue = NewSyncQueue(20)

func (h *ClientsManager) Run() {
	initActionQueue()
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	var signalsReceived uint
	for {
		select {
		case client := <-h.Register:
			h.Clients[client] = true
		case client := <-h.Unregister:
			if _, ok := h.Clients[client]; ok {
				delete(h.Clients, client)
				close(client.Send)
			}
		case message := <-models.ActionChan:
			LastActionsQueue.Push(message)
			for client := range h.Clients {
				select {
				case client.Send <- message:
				default:
					close(client.Send)
					delete(h.Clients, client)
				}
			}
		case s := <-sig:
			log.Info("received signal", s)
			signalsReceived++
			if signalsReceived < 2 {
				for client, _ := range h.Clients {
					delete(h.Clients, client)
					client.Close()
				}
				break

			}
		}
	}
}

func initActionQueue() {
	actions, err := models.GetLast20PublicFeeds()
	if err == nil {
		for i := len(actions) - 1; i >= 0; i-- {
			LastActionsQueue.Push(actions[i])
		}
	}
}