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, } }