oc-lib/tools/nats_caller.go

package tools

import (
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"cloud.o-forge.io/core/oc-lib/config"
	"cloud.o-forge.io/core/oc-lib/logs"
	"github.com/nats-io/nats.go"
	"github.com/rs/zerolog"
)

type NATSResponse struct {
	FromApp  string   `json:"from_app"`
	Datatype DataType `json:"datatype"`
	Method   int      `json:"method"`
	Payload  []byte   `json:"payload"`
}

// NATS Method Enum defines the different methods that can be used to interact with the NATS server
type NATSMethod int

var meths = []string{"remove execution", "create execution", "discovery", "workflow event", "peer discovery"}

const (
	REMOVE_EXECUTION NATSMethod = iota
	CREATE_EXECTUTION
	DISCOVERY
	WORKFLOW_EVENT

	REMOVE_PEER
	CREATE_PEER

	CREATE_RESOURCE
	REMOVE_RESOURCE
)

func (n NATSMethod) String() string {
	return meths[n]
}

// NameToMethod returns the NATSMethod enum value from a string
func NameToMethod(name string) NATSMethod {
	for _, v := range [...]NATSMethod{REMOVE_EXECUTION, CREATE_EXECTUTION, DISCOVERY, WORKFLOW_EVENT,
		REMOVE_PEER, CREATE_PEER, CREATE_RESOURCE, REMOVE_RESOURCE} {
		if strings.Contains(strings.ToLower(v.String()), strings.ToLower(name)) {
			return v
		}
	}
	return -1
}

// GenerateKey generates a key for the NATSMethod usefull for standard key based on data name & method
func (d NATSMethod) GenerateKey() string {
	return strings.ReplaceAll(d.String(), " ", "_")
}

type natsCaller struct{}

// NewNATSCaller creates a new instance of the NATS Caller
func NewNATSCaller() *natsCaller {
	return &natsCaller{}
}

// on workflows' scheduling. Messages must contain
// workflow execution ID, to allow retrieval of execution infos
func (s *natsCaller) ListenNats(execs map[NATSMethod]func(map[string]string)) {
	log := logs.GetLogger()
	if config.GetConfig().NATSUrl == "" {
		log.Error().Msg(" -> NATS_SERVER is not set")
		return
	}
	for {
		nc, err := nats.Connect(config.GetConfig().NATSUrl)
		if err != nil {
			log.Error().Msg("Could not connect to NATS")
			time.Sleep(1 * time.Minute)
			continue
		}
		defer nc.Close()
		var wg sync.WaitGroup
		wg.Add(len(execs))
		for k, v := range execs {
			go s.listenForChange(log, nc, k, v, &wg)
		}
		wg.Wait()
		break
	}
}

// SetNATSPub sets a message to the NATS server
func (o *natsCaller) SetNATSPub(method NATSMethod, data NATSResponse) string {
	if config.GetConfig().NATSUrl == "" {
		return " -> NATS_SERVER is not set"
	}
	for {
		nc, err := nats.Connect(config.GetConfig().NATSUrl)
		if err != nil {
			time.Sleep(1 * time.Minute)
			continue
		}
		defer nc.Close()
		js, err := json.Marshal(data)
		if err != nil {
			return " -> " + err.Error()
		}
		err = nc.Publish(method.GenerateKey(), js) // Publish the message on the NATS server with a channel name based on the data name (or whatever start) and the method
		if err != nil {
			time.Sleep(1 * time.Minute)
			continue
		}
		break
	}
	return ""
}

// Goroutine listening to a NATS server for updates
// on workflows' scheduling. Messages must contain
// workflow execution ID, to allow retrieval of execution infos
func (o *natsCaller) listenForChange(logger zerolog.Logger, nc *nats.Conn, natsTools NATSMethod,
	function func(map[string]string), wg *sync.WaitGroup) {
	defer wg.Done()
	ch := make(chan *nats.Msg, 64)
	logger.Info().Msg("Listening to " + natsTools.GenerateKey())
	subs, err := nc.ChanSubscribe(natsTools.GenerateKey(), ch)
	if err != nil {
		logger.Error().Msg("Error listening to NATS : " + err.Error())
	}
	defer subs.Unsubscribe()

	for msg := range ch {
		map_mess := map[string]string{}
		json.Unmarshal(msg.Data, &map_mess)
		fmt.Println("Catching " + natsTools.String() + " workflow... " + map_mess["id"])
		function(map_mess)
	}
}