versia-go/pkg/taskqueue/client.go

package taskqueue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"git.devminer.xyz/devminer/unitel"
	"github.com/go-logr/logr"
	"github.com/google/uuid"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

type taskWrapper struct {
	Task       Task              `json:"task"`
	EnqueuedAt time.Time         `json:"enqueuedAt"`
	TraceInfo  map[string]string `json:"traceInfo"`
}

func (c *Client) newTaskWrapper(ctx context.Context, task Task) taskWrapper {
	traceInfo := make(map[string]string)
	c.telemetry.InjectIntoMap(ctx, traceInfo)

	return taskWrapper{
		Task:       task,
		EnqueuedAt: time.Now(),
		TraceInfo:  traceInfo,
	}
}

type Task struct {
	ID      string
	Type    string
	Payload json.RawMessage
}

func NewTask(type_ string, payload any) (Task, error) {
	id := uuid.New()

	d, err := json.Marshal(payload)
	if err != nil {
		return Task{}, err
	}

	return Task{
		ID:      id.String(),
		Type:    type_,
		Payload: d,
	}, nil
}

type Client struct {
	name    string
	subject string

	nc *nats.Conn
	js jetstream.JetStream
	s  jetstream.Stream

	stopCh    chan struct{}
	closeOnce func()

	telemetry *unitel.Telemetry
	log       logr.Logger
}

func NewClient(streamName string, natsClient *nats.Conn, telemetry *unitel.Telemetry, log logr.Logger) (*Client, error) {
	js, err := jetstream.New(natsClient)
	if err != nil {
		return nil, err
	}

	return &Client{
		name: streamName,

		js: js,

		telemetry: telemetry,
		log:       log,
	}, nil
}

func (c *Client) Set(ctx context.Context, name string) (*Set, error) {
	streamName := fmt.Sprintf("%s_%s", c.name, name)

	s, err := c.js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              streamName,
		MaxConsumers:      -1,
		MaxMsgs:           -1,
		Discard:           jetstream.DiscardOld,
		MaxMsgsPerSubject: -1,
		Storage:           jetstream.FileStorage,
		Compression:       jetstream.S2Compression,
		AllowDirect:       true,
	})
	if errors.Is(err, nats.ErrStreamNameAlreadyInUse) {
		s, err = c.js.Stream(ctx, streamName)
		if err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	return &Set{
		handlers: make(map[string][]TaskHandler),

		streamName: streamName,
		c:          c,
		s:          s,
		log:        c.log.WithName(fmt.Sprintf("taskset(%s)", name)),
		telemetry:  c.telemetry,
	}, nil
}