kmsvc-sdk/messages.go

package kmsvc

import (
	"context"
	"fmt"

	kafkamgmtv1 "forgejo.riotpiao.com/rock/kmsvc-proto/gen/kafkamgmt/v1"
)

// MaxMessageBodyBytes matches SQS's message size cap (design.md §11.2).
const MaxMessageBodyBytes = 256 * 1024

// Message is a received message, decoupled from the generated protobuf type.
type Message struct {
	MessageID       string
	ReceiptHandle   string
	Body            []byte
	Attributes      map[string]string
	ReceiveCount    int32
	MessageGroupID  string
}

// SendMessageInput is the input to SendMessage.
type SendMessageInput struct {
	QueueName               string
	Body                    []byte
	Attributes              map[string]string
	MessageGroupID          string // FIFO only
	MessageDeduplicationID  string // FIFO only
	DelaySeconds            int32
}

// SendMessageOutput is the result of a successful SendMessage call.
type SendMessageOutput struct {
	MessageID      string
	SequenceNumber string // FIFO only
}

func checkBodySize(body []byte) error {
	if len(body) > MaxMessageBodyBytes {
		return fmt.Errorf("kmsvc: message body is %d bytes, exceeds %d byte limit: %w", len(body), MaxMessageBodyBytes, ErrMessageTooLarge)
	}
	return nil
}

// SendMessage sends a single message to queueName.
func (c *Client) SendMessage(ctx context.Context, in SendMessageInput) (*SendMessageOutput, error) {
	if err := checkBodySize(in.Body); err != nil {
		return nil, err
	}

	resp, err := c.stub.SendMessage(ctx, &kafkamgmtv1.SendMessageRequest{
		QueueName:              in.QueueName,
		MessageBody:            in.Body,
		MessageAttributes:      &kafkamgmtv1.MessageAttributes{Values: in.Attributes},
		MessageGroupId:         in.MessageGroupID,
		MessageDeduplicationId: in.MessageDeduplicationID,
		DelaySeconds:           in.DelaySeconds,
	})
	if err != nil {
		return nil, mapError(err)
	}

	return &SendMessageOutput{MessageID: resp.MessageId, SequenceNumber: resp.SequenceNumber}, nil
}

// SendMessageBatchEntry is one entry in a SendMessageBatch call.
type SendMessageBatchEntry struct {
	ID                     string
	Body                   []byte
	Attributes             map[string]string
	MessageGroupID         string
	MessageDeduplicationID string
	DelaySeconds           int32
}

// BatchResultEntry reports the outcome of one entry in a batch call.
type BatchResultEntry struct {
	ID        string
	MessageID string
	Error     string
}

// SendMessageBatchOutput is the result of a SendMessageBatch call.
type SendMessageBatchOutput struct {
	Successful []BatchResultEntry
	Failed     []BatchResultEntry
}

// SendMessageBatch sends up to 10 messages to queueName in one round trip.
func (c *Client) SendMessageBatch(ctx context.Context, queueName string, entries []SendMessageBatchEntry) (*SendMessageBatchOutput, error) {
	pbEntries := make([]*kafkamgmtv1.SendMessageBatchEntry, 0, len(entries))
	for _, e := range entries {
		if err := checkBodySize(e.Body); err != nil {
			return nil, fmt.Errorf("entry %q: %w", e.ID, err)
		}
		pbEntries = append(pbEntries, &kafkamgmtv1.SendMessageBatchEntry{
			Id:                     e.ID,
			MessageBody:            e.Body,
			MessageAttributes:      &kafkamgmtv1.MessageAttributes{Values: e.Attributes},
			MessageGroupId:         e.MessageGroupID,
			MessageDeduplicationId: e.MessageDeduplicationID,
			DelaySeconds:           e.DelaySeconds,
		})
	}

	resp, err := c.stub.SendMessageBatch(ctx, &kafkamgmtv1.SendMessageBatchRequest{
		QueueName: queueName,
		Entries:   pbEntries,
	})
	if err != nil {
		return nil, mapError(err)
	}

	return &SendMessageBatchOutput{
		Successful: toBatchResults(resp.Successful),
		Failed:     toBatchResults(resp.Failed),
	}, nil
}

func toBatchResults(in []*kafkamgmtv1.BatchResultEntry) []BatchResultEntry {
	out := make([]BatchResultEntry, 0, len(in))
	for _, e := range in {
		out = append(out, BatchResultEntry{ID: e.Id, MessageID: e.MessageId, Error: e.Error})
	}
	return out
}

// DeleteMessage acknowledges receiptHandle, removing the message permanently.
func (c *Client) DeleteMessage(ctx context.Context, queueName, receiptHandle string) error {
	_, err := c.stub.DeleteMessage(ctx, &kafkamgmtv1.DeleteMessageRequest{
		QueueName:     queueName,
		ReceiptHandle: receiptHandle,
	})
	if err != nil {
		return mapError(err)
	}
	return nil
}

// DeleteMessageBatchEntry is one entry in a DeleteMessageBatch call.
type DeleteMessageBatchEntry struct {
	ID            string
	ReceiptHandle string
}

// DeleteMessageBatchOutput is the result of a DeleteMessageBatch call.
type DeleteMessageBatchOutput struct {
	Successful []BatchResultEntry
	Failed     []BatchResultEntry
}

// DeleteMessageBatch acknowledges up to 10 messages in one round trip.
func (c *Client) DeleteMessageBatch(ctx context.Context, queueName string, entries []DeleteMessageBatchEntry) (*DeleteMessageBatchOutput, error) {
	pbEntries := make([]*kafkamgmtv1.DeleteMessageBatchEntry, 0, len(entries))
	for _, e := range entries {
		pbEntries = append(pbEntries, &kafkamgmtv1.DeleteMessageBatchEntry{
			Id:            e.ID,
			ReceiptHandle: e.ReceiptHandle,
		})
	}

	resp, err := c.stub.DeleteMessageBatch(ctx, &kafkamgmtv1.DeleteMessageBatchRequest{
		QueueName: queueName,
		Entries:   pbEntries,
	})
	if err != nil {
		return nil, mapError(err)
	}

	return &DeleteMessageBatchOutput{
		Successful: toBatchResults(resp.Successful),
		Failed:     toBatchResults(resp.Failed),
	}, nil
}

// ChangeMessageVisibility extends or shortens the visibility timeout for an
// in-flight message identified by receiptHandle.
func (c *Client) ChangeMessageVisibility(ctx context.Context, queueName, receiptHandle string, visibilityTimeoutSeconds int32) error {
	_, err := c.stub.ChangeMessageVisibility(ctx, &kafkamgmtv1.ChangeMessageVisibilityRequest{
		QueueName:                queueName,
		ReceiptHandle:            receiptHandle,
		VisibilityTimeoutSeconds: visibilityTimeoutSeconds,
	})
	if err != nil {
		return mapError(err)
	}
	return nil
}

func toMessage(m *kafkamgmtv1.Message) Message {
	var attrs map[string]string
	if m.Attributes != nil {
		attrs = m.Attributes.Values
	}
	return Message{
		MessageID:      m.MessageId,
		ReceiptHandle:  m.ReceiptHandle,
		Body:           m.Body,
		Attributes:     attrs,
		ReceiveCount:   m.ReceiveCount,
		MessageGroupID: m.MessageGroupId,
	}
}