From 2ceac063ded6dd7b8d0faa9e96f01ea8f702ab0f Mon Sep 17 00:00:00 2001 From: Nitin Kumar Date: Sun, 15 Mar 2026 22:40:59 +0530 Subject: [PATCH] add dead-letter queues, cost simulation, and serverless triggers - DLQ support for SQS, Service Bus, and Pub/Sub with configurable max receive count - New cost/ package for simulated cloud billing with default rates and custom overrides - Lambda/Function triggers on message queues across all 3 providers - 11 new tests (32 total), all passing --- cloudemu_test.go | 381 +++++++++++++++++++++++ cost/cost.go | 171 ++++++++++ messagequeue/driver/driver.go | 7 + providers/aws/sqs/sqs.go | 91 +++++- providers/azure/servicebus/servicebus.go | 90 +++++- providers/gcp/pubsub/pubsub.go | 90 +++++- 6 files changed, 815 insertions(+), 15 deletions(-) create mode 100644 cost/cost.go diff --git a/cloudemu_test.go b/cloudemu_test.go index 6fc09ef2..70f80c4f 100644 --- a/cloudemu_test.go +++ b/cloudemu_test.go @@ -3,12 +3,14 @@ package cloudemu import ( "context" "sort" + "sync/atomic" "testing" "time" "github.com/NitinKumar004/cloudemu/compute" computedriver "github.com/NitinKumar004/cloudemu/compute/driver" "github.com/NitinKumar004/cloudemu/config" + "github.com/NitinKumar004/cloudemu/cost" "github.com/NitinKumar004/cloudemu/database/driver" dnsdriver "github.com/NitinKumar004/cloudemu/dns/driver" cerrors "github.com/NitinKumar004/cloudemu/errors" @@ -20,6 +22,7 @@ import ( netdriver "github.com/NitinKumar004/cloudemu/networking/driver" "github.com/NitinKumar004/cloudemu/ratelimit" "github.com/NitinKumar004/cloudemu/recorder" + serverlessdriver "github.com/NitinKumar004/cloudemu/serverless/driver" "github.com/NitinKumar004/cloudemu/storage" storagedriver "github.com/NitinKumar004/cloudemu/storage/driver" ) @@ -1221,3 +1224,381 @@ func TestLifecycleStartEmitsRunningMetrics(t *testing.T) { t.Errorf("expected ALARM after StartInstances (CPU=25 > 0), got %s", alarms[0].State) } } + +// ============================================================================== +// Feature: Dead-Letter Queue Tests +// ============================================================================== + +func TestDeadLetterQueue(t *testing.T) { + ctx := context.Background() + clock := config.NewFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)) + p := NewAWS(config.WithClock(clock)) + + // 1. Create the DLQ first + dlq, err := p.SQS.CreateQueue(ctx, mqdriver.QueueConfig{ + Name: "my-queue-dlq", + }) + if err != nil { + t.Fatal(err) + } + + // 2. Create main queue with DLQ configured (maxReceiveCount=2) + mainQ, err := p.SQS.CreateQueue(ctx, mqdriver.QueueConfig{ + Name: "my-queue", + VisibilityTimeout: 1, + DeadLetterQueue: &mqdriver.DeadLetterConfig{ + TargetQueueURL: dlq.URL, + MaxReceiveCount: 2, + }, + }) + if err != nil { + t.Fatal(err) + } + + // 3. Send a message + _, err = p.SQS.SendMessage(ctx, mqdriver.SendMessageInput{ + QueueURL: mainQ.URL, + Body: "process me", + }) + if err != nil { + t.Fatal(err) + } + + // 4. Receive the message twice (simulating failed processing — not deleting it) + for i := 0; i < 2; i++ { + msgs, err := p.SQS.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{ + QueueURL: mainQ.URL, + }) + if err != nil { + t.Fatal(err) + } + if len(msgs) != 1 { + t.Fatalf("receive %d: expected 1 message, got %d", i+1, len(msgs)) + } + // Don't delete — simulating failure. Make it visible again. + clock.Advance(2 * time.Second) + } + + // 5. Third receive should trigger DLQ move (receiveCount exceeds maxReceiveCount=2) + msgs, err := p.SQS.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{ + QueueURL: mainQ.URL, + }) + if err != nil { + t.Fatal(err) + } + if len(msgs) != 0 { + t.Errorf("expected 0 messages in main queue after DLQ move, got %d", len(msgs)) + } + + // 6. Verify message is now in the DLQ + dlqInfo, err := p.SQS.GetQueueInfo(ctx, dlq.URL) + if err != nil { + t.Fatal(err) + } + if dlqInfo.ApproxMessageCount != 1 { + t.Errorf("expected 1 message in DLQ, got %d", dlqInfo.ApproxMessageCount) + } + + // 7. Receive from DLQ to verify message body + dlqMsgs, err := p.SQS.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{ + QueueURL: dlq.URL, + }) + if err != nil { + t.Fatal(err) + } + if len(dlqMsgs) != 1 { + t.Fatalf("expected 1 DLQ message, got %d", len(dlqMsgs)) + } + if dlqMsgs[0].Body != "process me" { + t.Errorf("expected DLQ message body 'process me', got %q", dlqMsgs[0].Body) + } +} + +func TestDeadLetterQueueAzure(t *testing.T) { + ctx := context.Background() + clock := config.NewFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)) + p := NewAzure(config.WithClock(clock)) + + dlq, err := p.ServiceBus.CreateQueue(ctx, mqdriver.QueueConfig{Name: "sb-dlq"}) + if err != nil { + t.Fatal(err) + } + + mainQ, err := p.ServiceBus.CreateQueue(ctx, mqdriver.QueueConfig{ + Name: "sb-main", + VisibilityTimeout: 1, + DeadLetterQueue: &mqdriver.DeadLetterConfig{TargetQueueURL: dlq.URL, MaxReceiveCount: 1}, + }) + if err != nil { + t.Fatal(err) + } + + p.ServiceBus.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: mainQ.URL, Body: "hello"}) + + // First receive succeeds + msgs, _ := p.ServiceBus.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: mainQ.URL}) + if len(msgs) != 1 { + t.Fatalf("expected 1 message, got %d", len(msgs)) + } + + clock.Advance(2 * time.Second) + + // Second receive moves to DLQ (receiveCount=2 > maxReceiveCount=1) + msgs, _ = p.ServiceBus.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: mainQ.URL}) + if len(msgs) != 0 { + t.Errorf("expected 0 messages after DLQ move, got %d", len(msgs)) + } + + dlqMsgs, _ := p.ServiceBus.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: dlq.URL}) + if len(dlqMsgs) != 1 || dlqMsgs[0].Body != "hello" { + t.Errorf("expected DLQ message with body 'hello', got %v", dlqMsgs) + } +} + +func TestDeadLetterQueueGCP(t *testing.T) { + ctx := context.Background() + clock := config.NewFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)) + p := NewGCP(config.WithClock(clock)) + + dlq, err := p.PubSub.CreateQueue(ctx, mqdriver.QueueConfig{Name: "ps-dlq"}) + if err != nil { + t.Fatal(err) + } + + mainQ, err := p.PubSub.CreateQueue(ctx, mqdriver.QueueConfig{ + Name: "ps-main", + VisibilityTimeout: 1, + DeadLetterQueue: &mqdriver.DeadLetterConfig{TargetQueueURL: dlq.URL, MaxReceiveCount: 1}, + }) + if err != nil { + t.Fatal(err) + } + + p.PubSub.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: mainQ.URL, Body: "gcp-msg"}) + + msgs, _ := p.PubSub.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: mainQ.URL}) + if len(msgs) != 1 { + t.Fatalf("expected 1 message, got %d", len(msgs)) + } + + clock.Advance(2 * time.Second) + + msgs, _ = p.PubSub.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: mainQ.URL}) + if len(msgs) != 0 { + t.Errorf("expected 0 messages after DLQ move, got %d", len(msgs)) + } + + dlqMsgs, _ := p.PubSub.ReceiveMessages(ctx, mqdriver.ReceiveMessageInput{QueueURL: dlq.URL}) + if len(dlqMsgs) != 1 || dlqMsgs[0].Body != "gcp-msg" { + t.Errorf("expected DLQ message with body 'gcp-msg', got %v", dlqMsgs) + } +} + +// ============================================================================== +// Feature: Cost Simulation Tests +// ============================================================================== + +func TestCostTracker(t *testing.T) { + tracker := cost.New() + + // Simulate some cloud operations + tracker.Record("compute", "RunInstances", 3) + tracker.Record("storage", "PutObject", 100) + tracker.Record("storage", "GetObject", 500) + tracker.Record("database", "PutItem", 1000) + tracker.Record("serverless", "Invoke", 10000) + tracker.Record("messagequeue", "SendMessage", 5000) + + // Verify total cost is > 0 + total := tracker.TotalCost() + if total <= 0 { + t.Errorf("expected total cost > 0, got %f", total) + } + + // Verify cost by service + byService := tracker.CostByService() + if byService["compute"] <= 0 { + t.Error("expected compute cost > 0") + } + if byService["storage"] <= 0 { + t.Error("expected storage cost > 0") + } + if byService["database"] <= 0 { + t.Error("expected database cost > 0") + } + + // Verify cost by operation + byOp := tracker.CostByOperation() + if byOp["compute:RunInstances"] <= 0 { + t.Error("expected RunInstances cost > 0") + } + + // Verify all costs recorded + allCosts := tracker.AllCosts() + if len(allCosts) != 6 { + t.Errorf("expected 6 cost records, got %d", len(allCosts)) + } +} + +func TestCostTrackerCustomRates(t *testing.T) { + tracker := cost.New() + + // Set custom rate + tracker.SetRate("compute", "RunInstances", 0.50) + + tracker.Record("compute", "RunInstances", 10) + + total := tracker.TotalCost() + expected := 5.0 // 0.50 * 10 + if total != expected { + t.Errorf("expected cost %f, got %f", expected, total) + } +} + +func TestCostTrackerReset(t *testing.T) { + tracker := cost.New() + + tracker.Record("compute", "RunInstances", 5) + if tracker.TotalCost() <= 0 { + t.Error("expected cost > 0 before reset") + } + + tracker.Reset() + if tracker.TotalCost() != 0 { + t.Errorf("expected 0 cost after reset, got %f", tracker.TotalCost()) + } +} + +// ============================================================================== +// Feature: Lambda-SQS Trigger Tests +// ============================================================================== + +func TestLambdaSQSTrigger(t *testing.T) { + ctx := context.Background() + p := NewAWS() + + // 1. Create a Lambda function + p.Lambda.RegisterHandler("processor", func(_ context.Context, payload []byte) ([]byte, error) { + return []byte("processed: " + string(payload)), nil + }) + _, err := p.Lambda.CreateFunction(ctx, serverlessdriver.FunctionConfig{ + Name: "processor", Runtime: "go1.x", Handler: "main", + }) + if err != nil { + t.Fatal(err) + } + + // 2. Create SQS queue + q, err := p.SQS.CreateQueue(ctx, mqdriver.QueueConfig{Name: "trigger-queue"}) + if err != nil { + t.Fatal(err) + } + + // 3. Wire the trigger: SQS → Lambda + var triggerCount int64 + p.SQS.SetTrigger(q.URL, func(queueURL string, msg mqdriver.Message) { + // Invoke lambda with the message body + _, invokeErr := p.Lambda.Invoke(ctx, serverlessdriver.InvokeInput{ + FunctionName: "processor", + Payload: []byte(msg.Body), + }) + if invokeErr != nil { + t.Errorf("trigger invoke failed: %v", invokeErr) + } + atomic.AddInt64(&triggerCount, 1) + }) + + // 4. Send messages — Lambda should be triggered automatically + for i := 0; i < 5; i++ { + _, err := p.SQS.SendMessage(ctx, mqdriver.SendMessageInput{ + QueueURL: q.URL, + Body: "message-body", + }) + if err != nil { + t.Fatal(err) + } + } + + // 5. Verify Lambda was triggered 5 times + if atomic.LoadInt64(&triggerCount) != 5 { + t.Errorf("expected 5 trigger invocations, got %d", triggerCount) + } +} + +func TestLambdaSQSTriggerRemove(t *testing.T) { + ctx := context.Background() + p := NewAWS() + + q, err := p.SQS.CreateQueue(ctx, mqdriver.QueueConfig{Name: "removable-trigger"}) + if err != nil { + t.Fatal(err) + } + + var triggerCount int64 + p.SQS.SetTrigger(q.URL, func(_ string, _ mqdriver.Message) { + atomic.AddInt64(&triggerCount, 1) + }) + + // Send one message — trigger fires + p.SQS.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "first"}) + if atomic.LoadInt64(&triggerCount) != 1 { + t.Errorf("expected 1 trigger, got %d", triggerCount) + } + + // Remove trigger + p.SQS.RemoveTrigger(q.URL) + + // Send another message — trigger should NOT fire + p.SQS.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "second"}) + if atomic.LoadInt64(&triggerCount) != 1 { + t.Errorf("expected still 1 trigger after removal, got %d", triggerCount) + } +} + +func TestAzureFunctionServiceBusTrigger(t *testing.T) { + ctx := context.Background() + p := NewAzure() + + q, err := p.ServiceBus.CreateQueue(ctx, mqdriver.QueueConfig{Name: "az-trigger-queue"}) + if err != nil { + t.Fatal(err) + } + + var received []string + p.ServiceBus.SetTrigger(q.URL, func(_ string, msg mqdriver.Message) { + received = append(received, msg.Body) + }) + + p.ServiceBus.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "azure-msg-1"}) + p.ServiceBus.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "azure-msg-2"}) + + if len(received) != 2 { + t.Errorf("expected 2 triggered messages, got %d", len(received)) + } + if received[0] != "azure-msg-1" || received[1] != "azure-msg-2" { + t.Errorf("unexpected messages: %v", received) + } +} + +func TestGCPCloudFunctionPubSubTrigger(t *testing.T) { + ctx := context.Background() + p := NewGCP() + + q, err := p.PubSub.CreateQueue(ctx, mqdriver.QueueConfig{Name: "gcp-trigger-topic"}) + if err != nil { + t.Fatal(err) + } + + var received []string + p.PubSub.SetTrigger(q.URL, func(_ string, msg mqdriver.Message) { + received = append(received, msg.Body) + }) + + p.PubSub.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "gcp-event-1"}) + p.PubSub.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "gcp-event-2"}) + p.PubSub.SendMessage(ctx, mqdriver.SendMessageInput{QueueURL: q.URL, Body: "gcp-event-3"}) + + if len(received) != 3 { + t.Errorf("expected 3 triggered messages, got %d", len(received)) + } +} diff --git a/cost/cost.go b/cost/cost.go new file mode 100644 index 00000000..af1c9851 --- /dev/null +++ b/cost/cost.go @@ -0,0 +1,171 @@ +// Package cost provides simulated cost tracking for cloud operations. +package cost + +import ( + "sync" +) + +// ServiceCost defines the cost structure for a single service operation. +type ServiceCost struct { + Service string + Operation string + UnitCost float64 + Quantity int + Total float64 +} + +// Tracker tracks simulated costs across all cloud operations. +type Tracker struct { + mu sync.RWMutex + costs []ServiceCost + rates map[string]float64 +} + +// New creates a new cost tracker with default cloud pricing rates. +func New() *Tracker { + return &Tracker{ + costs: make([]ServiceCost, 0), + rates: defaultRates(), + } +} + +// defaultRates returns approximate per-operation costs (simplified). +func defaultRates() map[string]float64 { + return map[string]float64{ + // Compute (per instance-hour) + "compute:RunInstances": 0.0116, // t2.micro equivalent + "compute:StartInstances": 0.0, + "compute:StopInstances": 0.0, + "compute:TerminateInstances": 0.0, + + // Storage (per operation) + "storage:PutObject": 0.000005, // $5 per 1M PUTs + "storage:GetObject": 0.0000004, + "storage:DeleteObject": 0.0, + "storage:ListObjects": 0.000005, + "storage:CreateBucket": 0.0, + + // Database (per operation) + "database:PutItem": 0.00000125, // 1 WCU + "database:GetItem": 0.00000025, // 0.5 RCU + "database:Query": 0.00000025, + "database:Scan": 0.00000025, + "database:BatchPutItems": 0.00000125, + "database:BatchGetItems": 0.00000025, + + // Serverless (per invocation) + "serverless:Invoke": 0.0000002, // $0.20 per 1M + + // Message Queue (per operation) + "messagequeue:SendMessage": 0.0000004, // $0.40 per 1M + "messagequeue:ReceiveMessages": 0.0000004, + + // DNS (per query) + "dns:CreateRecord": 0.0, + "dns:GetRecord": 0.0000004, + + // Monitoring (per metric) + "monitoring:PutMetricData": 0.00001, // $0.01 per 1K + "monitoring:GetMetricData": 0.00001, + + // Networking + "networking:CreateVPC": 0.0, + "networking:CreateSubnet": 0.0, + + // Load Balancer (per hour) + "loadbalancer:CreateLoadBalancer": 0.0225, + + // IAM (free) + "iam:CreateUser": 0.0, + "iam:CheckPermission": 0.0, + } +} + +// Record records a cost event for a service operation. +func (t *Tracker) Record(service, operation string, quantity int) { + t.mu.Lock() + defer t.mu.Unlock() + + key := service + ":" + operation + rate, ok := t.rates[key] + if !ok { + rate = 0.0 + } + + total := rate * float64(quantity) + + t.costs = append(t.costs, ServiceCost{ + Service: service, + Operation: operation, + UnitCost: rate, + Quantity: quantity, + Total: total, + }) +} + +// SetRate overrides the default rate for a service operation. +func (t *Tracker) SetRate(service, operation string, rate float64) { + t.mu.Lock() + defer t.mu.Unlock() + + t.rates[service+":"+operation] = rate +} + +// TotalCost returns the total simulated cost across all operations. +func (t *Tracker) TotalCost() float64 { + t.mu.RLock() + defer t.mu.RUnlock() + + var total float64 + for _, c := range t.costs { + total += c.Total + } + + return total +} + +// CostByService returns the total cost grouped by service. +func (t *Tracker) CostByService() map[string]float64 { + t.mu.RLock() + defer t.mu.RUnlock() + + result := make(map[string]float64) + for _, c := range t.costs { + result[c.Service] += c.Total + } + + return result +} + +// CostByOperation returns the total cost grouped by service:operation. +func (t *Tracker) CostByOperation() map[string]float64 { + t.mu.RLock() + defer t.mu.RUnlock() + + result := make(map[string]float64) + for _, c := range t.costs { + key := c.Service + ":" + c.Operation + result[key] += c.Total + } + + return result +} + +// AllCosts returns a copy of all recorded cost events. +func (t *Tracker) AllCosts() []ServiceCost { + t.mu.RLock() + defer t.mu.RUnlock() + + result := make([]ServiceCost, len(t.costs)) + copy(result, t.costs) + + return result +} + +// Reset clears all recorded costs. +func (t *Tracker) Reset() { + t.mu.Lock() + defer t.mu.Unlock() + + t.costs = t.costs[:0] +} diff --git a/messagequeue/driver/driver.go b/messagequeue/driver/driver.go index 58912f70..b4859a22 100644 --- a/messagequeue/driver/driver.go +++ b/messagequeue/driver/driver.go @@ -12,6 +12,13 @@ type QueueConfig struct { MaxMessageSize int MessageRetention int // seconds Tags map[string]string + DeadLetterQueue *DeadLetterConfig +} + +// DeadLetterConfig configures a dead-letter queue for failed messages. +type DeadLetterConfig struct { + TargetQueueURL string + MaxReceiveCount int // move to DLQ after this many receives } // QueueInfo describes a message queue. diff --git a/providers/aws/sqs/sqs.go b/providers/aws/sqs/sqs.go index d67d1e54..252c825e 100644 --- a/providers/aws/sqs/sqs.go +++ b/providers/aws/sqs/sqs.go @@ -28,6 +28,7 @@ type sqsMessage struct { ReceiptHandle string VisibleAt time.Time SentAt time.Time + ReceiveCount int } // queueData holds the internal state of a single SQS queue. @@ -41,22 +42,46 @@ type queueData struct { maxMessageSize int messageRetention int deduplicationIndex map[string]time.Time + dlqConfig *driver.DeadLetterConfig } +// LambdaTrigger is a function that gets called when a message is sent to a queue. +type LambdaTrigger func(queueURL string, message driver.Message) + // Mock is an in-memory mock implementation of the AWS SQS service. type Mock struct { - queues *memstore.Store[*queueData] - opts *config.Options + queues *memstore.Store[*queueData] + opts *config.Options + mu sync.RWMutex + triggers map[string]LambdaTrigger // queueURL → trigger } // New creates a new SQS mock with the given configuration options. func New(opts *config.Options) *Mock { return &Mock{ - queues: memstore.New[*queueData](), - opts: opts, + queues: memstore.New[*queueData](), + opts: opts, + triggers: make(map[string]LambdaTrigger), } } +// SetTrigger registers a Lambda trigger for a queue. When a message is sent to the +// queue, the trigger function is called automatically. +func (m *Mock) SetTrigger(queueURL string, fn LambdaTrigger) { + m.mu.Lock() + defer m.mu.Unlock() + + m.triggers[queueURL] = fn +} + +// RemoveTrigger removes a Lambda trigger from a queue. +func (m *Mock) RemoveTrigger(queueURL string) { + m.mu.Lock() + defer m.mu.Unlock() + + delete(m.triggers, queueURL) +} + // CreateQueue creates a new SQS queue. func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.QueueInfo, error) { if cfg.Name == "" { @@ -101,6 +126,7 @@ func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.Q maxMessageSize: cfg.MaxMessageSize, messageRetention: cfg.MessageRetention, deduplicationIndex: make(map[string]time.Time), + dlqConfig: cfg.DeadLetterQueue, } m.queues.Set(url, qd) @@ -224,6 +250,21 @@ func (m *Mock) SendMessage(_ context.Context, input driver.SendMessageInput) (*d qd.deduplicationIndex[input.DeduplicationID] = now } + // Fire Lambda trigger if registered. + m.mu.RLock() + trigger := m.triggers[input.QueueURL] + m.mu.RUnlock() + + if trigger != nil { + triggerMsg := driver.Message{ + MessageID: msgID, + Body: input.Body, + Attributes: attrs, + GroupID: input.GroupID, + } + trigger(input.QueueURL, triggerMsg) + } + return &driver.SendMessageOutput{ MessageID: msgID, }, nil @@ -256,7 +297,9 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp now := m.opts.Clock.Now() var results []driver.Message - for _, msg := range qd.messages { + var toRemove []int + + for i, msg := range qd.messages { if len(results) >= maxMessages { break } @@ -265,6 +308,16 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp continue } + msg.ReceiveCount++ + + // Check if message exceeded max receive count — move to DLQ. + if qd.dlqConfig != nil && qd.dlqConfig.MaxReceiveCount > 0 && msg.ReceiveCount > qd.dlqConfig.MaxReceiveCount { + m.moveToDLQ(qd.dlqConfig.TargetQueueURL, msg) + toRemove = append(toRemove, i) + + continue + } + // Generate a new receipt handle for this receive. receiptHandle := idgen.GenerateID("receipt-") msg.ReceiptHandle = receiptHandle @@ -284,6 +337,12 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp }) } + // Remove DLQ-moved messages in reverse order. + for i := len(toRemove) - 1; i >= 0; i-- { + idx := toRemove[i] + qd.messages = append(qd.messages[:idx], qd.messages[idx+1:]...) + } + if results == nil { results = []driver.Message{} } @@ -291,6 +350,28 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp return results, nil } +// moveToDLQ moves a message to the dead-letter queue. +func (m *Mock) moveToDLQ(dlqURL string, msg *sqsMessage) { + dlq, ok := m.queues.Get(dlqURL) + if !ok { + return + } + + dlq.mu.Lock() + defer dlq.mu.Unlock() + + dlqMsg := &sqsMessage{ + ID: msg.ID, + Body: msg.Body, + GroupID: msg.GroupID, + Attributes: msg.Attributes, + VisibleAt: m.opts.Clock.Now(), + SentAt: m.opts.Clock.Now(), + } + + dlq.messages = append(dlq.messages, dlqMsg) +} + // DeleteMessage deletes a message from the specified queue using its receipt handle. func (m *Mock) DeleteMessage(_ context.Context, queueURL, receiptHandle string) error { qd, ok := m.queues.Get(queueURL) diff --git a/providers/azure/servicebus/servicebus.go b/providers/azure/servicebus/servicebus.go index ecc4d64d..b9f4c685 100644 --- a/providers/azure/servicebus/servicebus.go +++ b/providers/azure/servicebus/servicebus.go @@ -28,6 +28,7 @@ type sbMessage struct { ReceiptHandle string VisibleAt time.Time SentAt time.Time + ReceiveCount int } // queueData holds the internal state of a single Service Bus queue. @@ -41,22 +42,45 @@ type queueData struct { maxMessageSize int messageRetention int deduplicationIndex map[string]time.Time + dlqConfig *driver.DeadLetterConfig } +// FunctionTrigger is a function that gets called when a message is sent to a queue. +type FunctionTrigger func(queueURL string, message driver.Message) + // Mock is an in-memory mock implementation of the Azure Service Bus service. type Mock struct { - queues *memstore.Store[*queueData] - opts *config.Options + queues *memstore.Store[*queueData] + opts *config.Options + mu sync.RWMutex + triggers map[string]FunctionTrigger // queueURL → trigger } // New creates a new Service Bus mock with the given configuration options. func New(opts *config.Options) *Mock { return &Mock{ - queues: memstore.New[*queueData](), - opts: opts, + queues: memstore.New[*queueData](), + opts: opts, + triggers: make(map[string]FunctionTrigger), } } +// SetTrigger registers an Azure Function trigger for a queue. +func (m *Mock) SetTrigger(queueURL string, fn FunctionTrigger) { + m.mu.Lock() + defer m.mu.Unlock() + + m.triggers[queueURL] = fn +} + +// RemoveTrigger removes a Function trigger from a queue. +func (m *Mock) RemoveTrigger(queueURL string) { + m.mu.Lock() + defer m.mu.Unlock() + + delete(m.triggers, queueURL) +} + // CreateQueue creates a new Service Bus queue. func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.QueueInfo, error) { if cfg.Name == "" { @@ -101,6 +125,7 @@ func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.Q maxMessageSize: cfg.MaxMessageSize, messageRetention: cfg.MessageRetention, deduplicationIndex: make(map[string]time.Time), + dlqConfig: cfg.DeadLetterQueue, } m.queues.Set(url, qd) @@ -223,6 +248,21 @@ func (m *Mock) SendMessage(_ context.Context, input driver.SendMessageInput) (*d qd.deduplicationIndex[input.DeduplicationID] = now } + // Fire Function trigger if registered. + m.mu.RLock() + trigger := m.triggers[input.QueueURL] + m.mu.RUnlock() + + if trigger != nil { + triggerMsg := driver.Message{ + MessageID: msgID, + Body: input.Body, + Attributes: attrs, + GroupID: input.GroupID, + } + trigger(input.QueueURL, triggerMsg) + } + return &driver.SendMessageOutput{ MessageID: msgID, }, nil @@ -255,7 +295,9 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp now := m.opts.Clock.Now() var results []driver.Message - for _, msg := range qd.messages { + var toRemove []int + + for i, msg := range qd.messages { if len(results) >= maxMessages { break } @@ -264,6 +306,16 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp continue } + msg.ReceiveCount++ + + // Check if message exceeded max receive count — move to DLQ. + if qd.dlqConfig != nil && qd.dlqConfig.MaxReceiveCount > 0 && msg.ReceiveCount > qd.dlqConfig.MaxReceiveCount { + m.moveToDLQ(qd.dlqConfig.TargetQueueURL, msg) + toRemove = append(toRemove, i) + + continue + } + // Generate a new receipt handle (lock token) for this receive. receiptHandle := idgen.GenerateID("sb-lock-") msg.ReceiptHandle = receiptHandle @@ -283,6 +335,12 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp }) } + // Remove DLQ-moved messages in reverse order. + for i := len(toRemove) - 1; i >= 0; i-- { + idx := toRemove[i] + qd.messages = append(qd.messages[:idx], qd.messages[idx+1:]...) + } + if results == nil { results = []driver.Message{} } @@ -290,6 +348,28 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp return results, nil } +// moveToDLQ moves a message to the dead-letter queue. +func (m *Mock) moveToDLQ(dlqURL string, msg *sbMessage) { + dlq, ok := m.queues.Get(dlqURL) + if !ok { + return + } + + dlq.mu.Lock() + defer dlq.mu.Unlock() + + dlqMsg := &sbMessage{ + ID: msg.ID, + Body: msg.Body, + GroupID: msg.GroupID, + Attributes: msg.Attributes, + VisibleAt: m.opts.Clock.Now(), + SentAt: m.opts.Clock.Now(), + } + + dlq.messages = append(dlq.messages, dlqMsg) +} + // DeleteMessage deletes (completes) a message from the specified queue using its receipt handle (lock token). func (m *Mock) DeleteMessage(_ context.Context, queueURL, receiptHandle string) error { qd, ok := m.queues.Get(queueURL) diff --git a/providers/gcp/pubsub/pubsub.go b/providers/gcp/pubsub/pubsub.go index 598a70f1..9f72e043 100644 --- a/providers/gcp/pubsub/pubsub.go +++ b/providers/gcp/pubsub/pubsub.go @@ -28,6 +28,7 @@ type pubsubMessage struct { ReceiptHandle string VisibleAt time.Time SentAt time.Time + ReceiveCount int } // queueData holds the internal state of a single Pub/Sub topic+subscription pair. @@ -41,22 +42,45 @@ type queueData struct { maxMessageSize int messageRetention int deduplicationIndex map[string]time.Time + dlqConfig *driver.DeadLetterConfig } +// FunctionTrigger is a function that gets called when a message is published to a topic. +type FunctionTrigger func(queueURL string, message driver.Message) + // Mock is an in-memory mock implementation of the GCP Pub/Sub service. type Mock struct { - queues *memstore.Store[*queueData] - opts *config.Options + queues *memstore.Store[*queueData] + opts *config.Options + mu sync.RWMutex + triggers map[string]FunctionTrigger // subscriptionURL → trigger } // New creates a new Pub/Sub mock with the given configuration options. func New(opts *config.Options) *Mock { return &Mock{ - queues: memstore.New[*queueData](), - opts: opts, + queues: memstore.New[*queueData](), + opts: opts, + triggers: make(map[string]FunctionTrigger), } } +// SetTrigger registers a Cloud Function trigger for a subscription. +func (m *Mock) SetTrigger(queueURL string, fn FunctionTrigger) { + m.mu.Lock() + defer m.mu.Unlock() + + m.triggers[queueURL] = fn +} + +// RemoveTrigger removes a Cloud Function trigger from a subscription. +func (m *Mock) RemoveTrigger(queueURL string) { + m.mu.Lock() + defer m.mu.Unlock() + + delete(m.triggers, queueURL) +} + // CreateQueue creates a new Pub/Sub topic and subscription pair. func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.QueueInfo, error) { if cfg.Name == "" { @@ -101,6 +125,7 @@ func (m *Mock) CreateQueue(_ context.Context, cfg driver.QueueConfig) (*driver.Q maxMessageSize: cfg.MaxMessageSize, messageRetention: cfg.MessageRetention, deduplicationIndex: make(map[string]time.Time), + dlqConfig: cfg.DeadLetterQueue, } m.queues.Set(url, qd) @@ -223,6 +248,21 @@ func (m *Mock) SendMessage(_ context.Context, input driver.SendMessageInput) (*d qd.deduplicationIndex[input.DeduplicationID] = now } + // Fire Cloud Function trigger if registered. + m.mu.RLock() + trigger := m.triggers[input.QueueURL] + m.mu.RUnlock() + + if trigger != nil { + triggerMsg := driver.Message{ + MessageID: msgID, + Body: input.Body, + Attributes: attrs, + GroupID: input.GroupID, + } + trigger(input.QueueURL, triggerMsg) + } + return &driver.SendMessageOutput{ MessageID: msgID, }, nil @@ -255,7 +295,9 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp now := m.opts.Clock.Now() var results []driver.Message - for _, msg := range qd.messages { + var toRemove []int + + for i, msg := range qd.messages { if len(results) >= maxMessages { break } @@ -264,6 +306,16 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp continue } + msg.ReceiveCount++ + + // Check if message exceeded max receive count — move to DLQ. + if qd.dlqConfig != nil && qd.dlqConfig.MaxReceiveCount > 0 && msg.ReceiveCount > qd.dlqConfig.MaxReceiveCount { + m.moveToDLQ(qd.dlqConfig.TargetQueueURL, msg) + toRemove = append(toRemove, i) + + continue + } + // Generate a new receipt handle (ack ID) for this pull. receiptHandle := idgen.GenerateID("ack-") msg.ReceiptHandle = receiptHandle @@ -283,6 +335,12 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp }) } + // Remove DLQ-moved messages in reverse order. + for i := len(toRemove) - 1; i >= 0; i-- { + idx := toRemove[i] + qd.messages = append(qd.messages[:idx], qd.messages[idx+1:]...) + } + if results == nil { results = []driver.Message{} } @@ -290,6 +348,28 @@ func (m *Mock) ReceiveMessages(_ context.Context, input driver.ReceiveMessageInp return results, nil } +// moveToDLQ moves a message to the dead-letter queue. +func (m *Mock) moveToDLQ(dlqURL string, msg *pubsubMessage) { + dlq, ok := m.queues.Get(dlqURL) + if !ok { + return + } + + dlq.mu.Lock() + defer dlq.mu.Unlock() + + dlqMsg := &pubsubMessage{ + ID: msg.ID, + Body: msg.Body, + GroupID: msg.GroupID, + Attributes: msg.Attributes, + VisibleAt: m.opts.Clock.Now(), + SentAt: m.opts.Clock.Now(), + } + + dlq.messages = append(dlq.messages, dlqMsg) +} + // DeleteMessage acknowledges and removes a message from the subscription using its ack ID (receipt handle). func (m *Mock) DeleteMessage(_ context.Context, queueURL, receiptHandle string) error { qd, ok := m.queues.Get(queueURL)