LAPP_AdvancedScheduler/messaging/Contexts.go

package messaging

import (
	"github.com/kataras/neffos"
	"sync"
)

type Contexts struct {
	locker                 sync.RWMutex
	connectionOfContexts   map[string]*neffos.NSConn
	connectionId2ContextID map[string]string
	callbacksOfAdd         []func(string, *neffos.NSConn)
	callbacksOfRemove      []func(string)
	authenticator          func(map[string]string) (string, error)
}

func newContexts() *Contexts {
	return &Contexts{
		connectionOfContexts:   make(map[string]*neffos.NSConn, 10),
		connectionId2ContextID: make(map[string]string, 10),
		callbacksOfAdd:         make([]func(string, *neffos.NSConn), 0, 10),
		callbacksOfRemove:      make([]func(string), 0, 10),
	}
}

func (contexts *Contexts) add(contextID string, connection *neffos.NSConn) {
	locker.Lock()
	defer locker.Unlock()
	contexts.connectionOfContexts[contextID] = connection
	contexts.connectionId2ContextID[connection.Conn.ID()] = contextID
	for _, callback := range contexts.callbacksOfAdd {
		callback(contextID, connection)
	}
}

func (contexts *Contexts) remove(connectionId string) {
	locker.Lock()
	defer locker.Unlock()
	contextID, has := contexts.connectionId2ContextID[connectionId]
	if !has {
		return
	}
	delete(contexts.connectionId2ContextID, connectionId)
	delete(contexts.connectionOfContexts, contextID)
	for _, callback := range contexts.callbacksOfRemove {
		callback(contextID)
	}
}

func (contexts *Contexts) Get(contextID string) *neffos.NSConn {
	locker.RLock()
	defer locker.RUnlock()
	connection, ok := contexts.connectionOfContexts[contextID]
	if ok {
		return connection
	} else {
		return nil
	}
}

func (contexts *Contexts) OnAdd(callback func(string, *neffos.NSConn)) {
	locker.Lock()
	defer locker.Unlock()
	contexts.callbacksOfAdd = append(contexts.callbacksOfAdd, callback)
}

func (contexts *Contexts) OnRemove(callback func(string)) {
	locker.Lock()
	defer locker.Unlock()
	contexts.callbacksOfRemove = append(contexts.callbacksOfRemove, callback)
}

func (contexts *Contexts) SetAuthenticator(authenticator func(map[string]string) (string, error)) {
	contexts.authenticator = authenticator
}