queue.go

/*
Copyright 2024 Blnk Finance Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package blnk

import (
	"context"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"time"

	"github.com/blnkfinance/blnk/config"
	"github.com/blnkfinance/blnk/internal/hotpairs"
	"github.com/blnkfinance/blnk/internal/metrics"
	redis_db "github.com/blnkfinance/blnk/internal/redis-db"
	"github.com/sirupsen/logrus"
	"go.opentelemetry.io/otel/attribute"
	otelmetric "go.opentelemetry.io/otel/metric"

	"github.com/blnkfinance/blnk/model"
	"github.com/hibiken/asynq"
)

// Queue represents a queue for handling various tasks.
type Queue struct {
	Client    *asynq.Client
	Inspector *asynq.Inspector
	config    *config.Configuration
}

// TransactionTypePayload represents the payload for a transaction type.
type TransactionTypePayload struct {
	Data model.Transaction
}

// NewQueue initializes a new Queue instance with the provided configuration.
//
// Parameters:
// - conf *config.Configuration: The configuration for the queue.
//
// Returns:
// - *Queue: A pointer to the newly created Queue instance.
func NewQueue(conf *config.Configuration, client *asynq.Client) *Queue {
	redisOption, err := redis_db.ParseRedisURL(conf.Redis.Dns, conf.Redis.SkipTLSVerify)
	if err != nil {
		logrus.WithError(err).Fatal("failed to parse Redis URL")
	}

	queueOptions := asynq.RedisClientOpt{Addr: redisOption.Addr, Password: redisOption.Password, DB: redisOption.DB, TLSConfig: redisOption.TLSConfig, PoolSize: conf.Redis.PoolSize}
	inspector := asynq.NewInspector(queueOptions)
	return &Queue{
		Client:    client,
		Inspector: inspector,
		config:    conf,
	}
}

// queueInflightExpiry enqueues a task to handle inflight expiry for a transaction.
//
// Parameters:
// - transactionID string: The ID of the transaction.
// - expiresAt time.Time: The expiration time for the inflight status.
//
// Returns:
// - error: An error if the task could not be enqueued.
func (q *Queue) queueInflightExpiry(transactionID string, expiresAt time.Time) error {
	IPayload, err := json.Marshal(transactionID)
	if err != nil {
		return err
	}
	taskOptions := []asynq.Option{
		asynq.TaskID(transactionID),
		asynq.Queue(q.config.Queue.InflightExpiryQueue),
		asynq.ProcessIn(time.Until(expiresAt)),
	}
	task := asynq.NewTask(q.config.Queue.InflightExpiryQueue, IPayload, taskOptions...)
	_, err = q.Client.Enqueue(task)
	if err != nil {
		logrus.WithError(err).WithField("transaction_id", transactionID).Error("failed to enqueue inflight expiry")
		return err
	}
	logrus.WithField("transaction_id", transactionID).Debug("successfully enqueued inflight expiry")
	return nil
}

// queueIndexBatch enqueues a batch of items to be indexed in dependency order.
// This ensures that dependencies (e.g., balances) are indexed before the primary item (e.g., transaction).
// Uses the same IndexQueue but with a different task type for routing.
//
// Parameters:
// - batch interface{}: The batch containing dependencies and primary item to index.
//
// Returns:
// - error: An error if the task could not be enqueued.
func (q *Queue) queueIndexBatch(batch interface{}) error {
	if q.config.TypeSense.Dns == "" {
		return nil
	}

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

	taskOptions := []asynq.Option{asynq.Queue(q.config.Queue.IndexQueue)}
	task := asynq.NewTask("new:index:batch", payload, taskOptions...)
	_, err = q.Client.Enqueue(task)
	if err != nil {
		logrus.WithError(err).Error("failed to enqueue index batch")
		return err
	}
	logrus.Debug("successfully enqueued index batch")
	return nil
}

// queueIndexData enqueues a task to index data in a specified collection.
//
// Parameters:
// - id string: The ID of the data to index.
// - collection string: The name of the collection to index the data in.
// - data interface{}: The data to be indexed.
//
// Returns:
// - error: An error if the task could not be enqueued.
func (q *Queue) queueIndexData(id string, collection string, data interface{}) error {
	if q.config.TypeSense.Dns == "" {
		return nil
	}

	payload := map[string]interface{}{
		"collection": collection,
		"payload":    data,
	}

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

	taskOptions := []asynq.Option{asynq.Queue(q.config.Queue.IndexQueue)}
	task := asynq.NewTask(q.config.Queue.IndexQueue, IPayload, taskOptions...)
	_, err = q.Client.Enqueue(task)
	if err != nil {
		logrus.WithError(err).WithField("id", id).Error("failed to enqueue index data")
		return err
	}
	logrus.WithField("id", id).Debug("successfully enqueued index data")
	return nil
}

// Enqueue enqueues a transaction to the Redis queue.
//
// Parameters:
// - ctx context.Context: The context for the operation.
// - transaction *model.Transaction: The transaction to be enqueued.
//
// Returns:
// - error: An error if the transaction could not be enqueued.
func (q *Queue) Enqueue(ctx context.Context, transaction *model.Transaction) error {
	ctx, span := tracer.Start(ctx, "Adding Transaction To Redis Queue")
	defer span.End()

	payload, err := json.Marshal(transaction)
	if err != nil {
		return err
	}
	task := q.geTask(transaction, payload)
	_, err = q.Client.EnqueueContext(ctx, task, asynq.MaxRetry(q.config.Queue.MaxRetryAttempts))
	if err != nil {
		logrus.WithError(err).WithField("reference", transaction.Reference).Error("failed to enqueue transaction")
		return err
	}
	logrus.WithField("reference", transaction.Reference).Debug("successfully enqueued transaction")

	// Record enqueue metrics.
	metrics.QueueEnqueuedTotal.Add(ctx, 1,
		otelmetric.WithAttributes(attribute.String("queue_name", task.Type())),
	)

	return nil
}

// QueueInflightExpiry handles queuing a transaction for inflight expiration.
// This method is separate from the main Enqueue to ensure expiration is handled
// regardless of whether the transaction is queued or processed immediately.
//
// Parameters:
// - ctx context.Context: The context for the operation.
// - transaction *model.Transaction: The transaction to queue for expiration.
//
// Returns:
// - error: An error if the expiration could not be queued.
func (q *Queue) QueueInflightExpiry(ctx context.Context, transaction *model.Transaction) error {
	if !transaction.InflightExpiryDate.IsZero() {
		return q.queueInflightExpiry(transaction.TransactionID, transaction.InflightExpiryDate)
	}
	return nil
}

// geTask generates a task for a transaction and assigns it to a specific queue based on the balance ID.
// It ensures that transactions are evenly distributed across multiple queues by hashing the balance ID.
// This approach helps to avoid race conditions on a balance by ensuring that all transactions related to the same balance
// are processed serially within the same queue, thereby maintaining accuracy and consistency.
//
// Parameters:
// - transaction *model.Transaction: The transaction for which to generate the task.
// - payload []byte: The payload for the task, typically the serialized transaction data.
//
// Returns:
// - *asynq.Task: The generated task ready to be enqueued.
func (q *Queue) geTask(transaction *model.Transaction, payload []byte) *asynq.Task {
	queueName := q.transactionQueueName(transaction)

	taskOptions := []asynq.Option{asynq.TaskID(transaction.TransactionID), asynq.Queue(queueName)}
	if !transaction.ScheduledFor.IsZero() {
		taskOptions = append(taskOptions, asynq.ProcessIn(time.Until(transaction.ScheduledFor)))
	}

	return asynq.NewTask(queueName, payload, taskOptions...)
}

func (q *Queue) transactionQueueName(transaction *model.Transaction) string {
	if q.config.Queue.EnableHotLane && transaction != nil && transaction.MetaData != nil {
		if hotpairs.QueueLaneFromMetadata(transaction.MetaData) == hotpairs.LaneHot {
			metrics.HotpairsLaneRoutedTotal.Add(context.Background(), 1,
				otelmetric.WithAttributes(attribute.String("lane", "hot")),
			)
			return q.config.Queue.HotQueueName
		}
	}

	metrics.HotpairsLaneRoutedTotal.Add(context.Background(), 1,
		otelmetric.WithAttributes(attribute.String("lane", "normal")),
	)
	queueIndex := hashBalanceID(transaction.Source) % q.config.Queue.NumberOfQueues
	return fmt.Sprintf("%s_%d", q.config.Queue.TransactionQueue, queueIndex+1)
}

// hashBalanceID returns a consistent hash value for a string balance ID.
//
// Parameters:
// - balanceID string: The balance ID to hash.
//
// Returns:
// - int: The hash value of the balance ID.
func hashBalanceID(balanceID string) int {
	hasher := fnv.New32a()
	_, _ = hasher.Write([]byte(balanceID))
	return int(hasher.Sum32())
}

// GetTransactionFromQueue retrieves a transaction from the queue by its ID.
//
// Parameters:
// - transactionID string: The ID of the transaction to retrieve.
//
// Returns:
// - *model.Transaction: A pointer to the Transaction model if found.
// - error: An error if the transaction could not be retrieved.
func (q *Queue) GetTransactionFromQueue(transactionID string) (*model.Transaction, error) {
	for i := 1; i <= q.config.Queue.NumberOfQueues; i++ {
		queueName := fmt.Sprintf("%s_%d", q.config.Queue.TransactionQueue, i)
		task, err := q.Inspector.GetTaskInfo(queueName, transactionID)
		if err == nil && task != nil {
			var txn model.Transaction
			if err := json.Unmarshal(task.Payload, &txn); err != nil {
				return nil, err
			}
			return &txn, nil
		}
	}
	if q.config.Queue.EnableHotLane {
		task, err := q.Inspector.GetTaskInfo(q.config.Queue.HotQueueName, transactionID)
		if err == nil && task != nil {
			var txn model.Transaction
			if err := json.Unmarshal(task.Payload, &txn); err != nil {
				return nil, err
			}
			return &txn, nil
		}
	}
	return nil, nil // Return nil if transaction is not found in any queue
}

// queueInflightCommit enqueues a task to handle inflight commit for a transaction.
//
// Parameters:
// - transactionID string: The ID of the transaction.
// - commitAt time.Time: The scheduled time to automatically commit the inflight transaction.
//
// Returns:
// - error: An error if the task could not be enqueued.
func (q *Queue) queueInflightCommit(transactionID string, commitAt time.Time) error {
	IPayload, err := json.Marshal(transactionID)
	if err != nil {
		return err
	}

	taskOptions := []asynq.Option{
		asynq.TaskID(transactionID),
		asynq.Queue(q.config.Queue.InflightCommitQueue),
		asynq.ProcessIn(time.Until(commitAt)),
	}

	task := asynq.NewTask(q.config.Queue.InflightCommitQueue, IPayload, taskOptions...)
	_, err = q.Client.Enqueue(task)
	if err != nil {
		logrus.WithError(err).WithField("transaction_id", transactionID).Error("failed to enqueue inflight commit")
		return err
	}

	logrus.WithField("transaction_id", transactionID).Debug("successfully enqueued inflight commit")
	return nil
}

// QueueInflightCommit schedules an automatic commit for an inflight transaction at the specified date.
//
// Parameters:
// - ctx context.Context: The context for the operation.
// - transaction *model.Transaction: The transaction to be committed automatically.
//
// Returns:
// - error: An error if the task could not be enqueued.
func (q *Queue) QueueInflightCommit(ctx context.Context, transaction *model.Transaction) error {
	if !transaction.InflightCommitDate.IsZero() {
		return q.queueInflightCommit(transaction.TransactionID, transaction.InflightCommitDate)
	}
	return nil
}