aiforge/services/socketwrap/clientManager.go

package socketwrap

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

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"github.com/elliotchance/orderedmap"
)

type ClientsManager struct {
	Clients    *orderedmap.OrderedMap
	Register   chan *Client
	Unregister chan *Client
}

func NewClientsManager() *ClientsManager {
	return &ClientsManager{
		Register:   make(chan *Client),
		Unregister: make(chan *Client),
		Clients:    orderedmap.NewOrderedMap(),
	}
}

const MaxClients = 100

var LastActionsQueue = NewSyncQueue(15)

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.Set(client, true)
			if h.Clients.Len() > MaxClients {
				h.Clients.Delete(h.Clients.Front().Key)
			}

		case client := <-h.Unregister:
			if _, ok := h.Clients.Get(client); ok {
				h.Clients.Delete(client)
				close(client.Send)
			}
		case message := <-models.ActionChan:
			LastActionsQueue.Push(message)
			for _, client := range h.Clients.Keys() {
				select {
				case client.(*Client).Send <- message:
				default:
					close(client.(*Client).Send)
					h.Clients.Delete(client)
				}
			}
		case s := <-sig:
			log.Info("received signal", s)
			signalsReceived++
			if signalsReceived < 2 {
				for _, client := range h.Clients.Keys() {
					h.Clients.Delete(client)
					client.(*Client).Close()
				}
				break

			}
		}
	}
}

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

			user, err := models.GetUserByID(actions[i].UserID)
			if err == nil {
				if !user.IsOrganization() {
					LastActionsQueue.Push(actions[i])
				}

			}

		}
	}
}