diff --git a/README.md b/README.md index a24f35b..8b046c9 100644 --- a/README.md +++ b/README.md @@ -37,8 +37,40 @@ func main() { ``` ## Description -The firequeue utilizes an in-memory queue to ensure input to Amazon Kinesis Data Firehose. -When the looping process is cancelled by the context, the firequeue wait for the queue to be empty and then exit the loop. +firequeue owns an in-memory queue. `Enqueue` copies payloads; acceptance does not +guarantee persistence. Stop producers before cancelling `Loop`'s context. The queue +then rejects new input and drains with a fresh context, for at most 10 seconds by +default (`DrainTimeout` changes this). Always inspect `Loop`'s returned error: a +timeout reports the remaining event and payload-byte counts. Process loss can lose +pending data; ambiguous request outcomes can duplicate records. + +By default each input remains a separate record. For newline-delimited JSON, callers +can opt into concatenation without re-encoding or adding delimiters: + +```go +fq := firequeue.New(fh, "device-messages", + firequeue.Aggregate(64*1024, 500*time.Millisecond), +) +``` + +The first event sets the aggregation deadline; later arrivals do not extend it. +Payloads larger than the target remain standalone. Delivery backlog and retries +can add latency after sealing. Each payload must contain 1–1,024,000 bytes. Each +request contains at most 500 records and 4,000,000 payload bytes, before Base64. + +`MaxQueueLength` (default 100,000 events) and `MaxQueueBytes` (default 64 MiB of +payload, excluding allocation overhead) include aggregation, ready, retry, and +in-flight data. Capacity rejection is explicit. HTTP-success partial failures retry +only failed records with the same bytes. Invalid responses and ambiguous transport +failures retry all affected records. Retryable requests use exponential backoff. + +`Stats` separates input events from attempted/acknowledged Firehose records. +`AcceptedEvents = SuccessfulEvents + FailedEvents + PendingEvents`. +Terminal API failures invoke `ErrorHandler` and increment `FailedEvents`; `Loop` +can still return nil once the queue is empty. Estimated billable bytes round each +attempt up to 5 KiB and do not represent an AWS invoice or SDK-internal retries. +`AggregationWaitObserver` measures creation-to-seal time once per record, when the +first send attempt starts; callbacks must be concurrency-safe and return promptly. ## Installation diff --git a/delivery.go b/delivery.go new file mode 100644 index 0000000..316ac4a --- /dev/null +++ b/delivery.go @@ -0,0 +1,130 @@ +package firequeue + +import ( + "context" + "fmt" + "log" + "strings" + "time" + + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/firehose" +) + +const requestTimeout = 30 * time.Second + +func (q *Queue) send(ctx context.Context, batch []*queuedRecord) bool { + records := make([]*firehose.Record, len(batch)) + var billableBytes int64 + for i, record := range batch { + records[i] = record.record + billableBytes += int64((len(record.record.Data) + billingUnitBytes - 1) / billingUnitBytes * billingUnitBytes) + } + q.mu.Lock() + q.stats.BatchLength = int64(len(batch)) + q.stats.SentRecords += int64(len(batch)) + q.stats.EstimatedBillableBytes += billableBytes + q.mu.Unlock() + requestCtx, cancel := context.WithTimeout(ctx, requestTimeout) + defer cancel() + resp, err := q.firehose.PutRecordBatchWithContext(requestCtx, &firehose.PutRecordBatchInput{ + DeliveryStreamName: q.DeliveryStreamName, + Records: records, + }) + if err != nil { + if ctx.Err() != nil || requestCtx.Err() != nil || isRetryable(err) { + q.retry(batch) + if ctx.Err() == nil { + q.handleError(err, records) + } + return true + } + q.finish(batch, false) + q.handleError(err, records) + return false + } + if err := validateResponse(resp, len(batch)); err != nil { + q.retry(batch) + q.handleError(err, records) + return true + } + var succeeded, failed []*queuedRecord + for i, record := range batch { + if resp.RequestResponses[i].ErrorCode != nil && *resp.RequestResponses[i].ErrorCode != "" { + failed = append(failed, record) + } else { + succeeded = append(succeeded, record) + } + } + q.finish(succeeded, true) + q.retry(failed) + return len(failed) > 0 +} + +func validateResponse(resp *firehose.PutRecordBatchOutput, count int) error { + if resp == nil || resp.FailedPutCount == nil || len(resp.RequestResponses) != count { + return fmt.Errorf("invalid PutRecordBatch response for %d records", count) + } + var failed int64 + for _, entry := range resp.RequestResponses { + if entry == nil { + return fmt.Errorf("nil PutRecordBatch response entry") + } + if entry.ErrorCode != nil && *entry.ErrorCode != "" { + failed++ + } else if entry.RecordId == nil || *entry.RecordId == "" { + return fmt.Errorf("missing PutRecordBatch acknowledgment ID") + } + } + if failed != *resp.FailedPutCount { + return fmt.Errorf("inconsistent PutRecordBatch response: %d failures, declared %d", failed, *resp.FailedPutCount) + } + return nil +} + +func (q *Queue) finish(batch []*queuedRecord, succeeded bool) { + q.mu.Lock() + defer q.mu.Unlock() + for _, record := range batch { + delete(q.inFlight, record) + q.stats.PendingEvents -= record.events + q.stats.PendingBytes -= int64(len(record.record.Data)) + if succeeded { + q.stats.Success++ + q.stats.SuccessfulEvents += record.events + } else { + q.stats.UnretryableError++ + q.stats.FailedEvents += record.events + } + } +} + +func (q *Queue) retry(batch []*queuedRecord) { + q.mu.Lock() + defer q.mu.Unlock() + for _, record := range batch { + delete(q.inFlight, record) + q.stats.RetryCount++ + q.queue = append(q.queue, record) + } +} + +func (q *Queue) handleError(err error, records []*firehose.Record) { + if q.errorHandler != nil { + q.errorHandler(err, records) + return + } + log.Println(err) +} + +func isRetryable(err error) bool { + if failure, ok := err.(awserr.RequestFailure); ok { + switch failure.StatusCode() { + case 429, 500, 502, 503, 504: + return true + } + } + return request.IsErrorRetryable(err) || request.IsErrorThrottle(err) || + strings.Contains(err.Error(), "read: connection reset by peer") +} diff --git a/firequeue.go b/firequeue.go index 24eaf73..8bf4eb6 100644 --- a/firequeue.go +++ b/firequeue.go @@ -3,86 +3,119 @@ package firequeue import ( "context" "errors" - "expvar" "fmt" - "log" - "strings" "sync" - "sync/atomic" "time" "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/service/firehose" "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/cenkalti/backoff/v4" ) -// Option is a type for constructor options +const ( + maxRecordBytes = 1000 * 1024 + maxBatchBytes = 4000000 + maxBatchRecords = 500 + billingUnitBytes = 5 * 1024 + defaultQueueLength = 100000 + defaultQueueBytes = 64 * 1024 * 1024 + defaultBatchInterval = 2 * time.Millisecond + defaultDrainTimeout = 10 * time.Second + startTimeout = 2 * time.Second + idleInterval = time.Hour +) + +// Option configures a Queue before Loop starts. type Option func(*Queue) -// Parallel is an option to specify the number of parallelism -func Parallel(i int) Option { - return func(q *Queue) { - q.para = i +// Parallel sets the number of concurrent delivery workers. +func Parallel(n int) Option { + if n <= 0 { + panic("parallelism must be positive") } + return func(q *Queue) { q.para = n } } -// ErrorHandler is an option to specify the error handler +// ErrorHandler receives API errors and their affected records. +// It must be concurrency-safe, return promptly, and leave record data unchanged. func ErrorHandler(fn func(error, []*firehose.Record)) Option { - return func(q *Queue) { - q.errorHandler = fn - } + return func(q *Queue) { q.errorHandler = fn } } -// MaxQueueLength is an option to specify max length of in-memory queue -func MaxQueueLength(length int) Option { - if length <= 0 { +// MaxQueueLength bounds pending input events, including in-flight events. +func MaxQueueLength(n int) Option { + if n <= 0 { panic("max queue length must be positive") } - return func(q *Queue) { - q.maxQueueLength = length - } + return func(q *Queue) { q.maxQueueLength = n } } -// BatchSize is an option to specify batch size -func BatchSize(length int) Option { - if length <= 0 { - panic("max queue length must be positive") +// MaxQueueBytes bounds pending payload bytes, including in-flight and retry data. +func MaxQueueBytes(n int) Option { + if n <= 0 { + panic("max queue bytes must be positive") } - // "Each PutRecordBatch request supports up to 500 records." - // ref. https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/firehose#Client.PutRecordBatch - if length > 500 { - panic("max queue length must be less than 500") + return func(q *Queue) { q.maxQueueBytes = n } +} + +// BatchSize bounds Firehose records per request. +func BatchSize(n int) Option { + if n <= 0 || n > maxBatchRecords { + panic("batch size must be between 1 and 500") } - return func(q *Queue) { - q.batchSize = length + return func(q *Queue) { q.batchSize = n } +} + +// BatchInterval controls the delay between successful requests with queued data. +func BatchInterval(interval time.Duration) Option { + if interval <= time.Microsecond { + panic("batch interval must exceed one microsecond") } + return func(q *Queue) { q.batchInterval = interval } } -// BatchInterval is an option to specify batch interval -func BatchInterval(period time.Duration) Option { - // API limit is 1000 requests per second per shard(default). - // This means an average of 1 request per millisecond. - // ref. https://docs.aws.amazon.com/ja_jp/firehose/latest/dev/limits.html#:~:text=Direct%20PUT%20%E3%81%8C%E3%83%87%E3%83%BC%E3%82%BF%E3%82%BD%E3%83%BC%E3%82%B9%E3%81%A8%E3%81%97%E3%81%A6%E8%A8%AD%E5%AE%9A%E3%81%95%E3%82%8C%E3%81%A6%E3%81%84%E3%82%8B%E5%A0%B4%E5%90%88%E3%80%81%E5%90%84%20Kinesis%20Data%20Firehose%20%E3%83%87%E3%83%AA%E3%83%90%E3%83%AA%E3%83%BC%E3%82%B9%E3%83%88%E3%83%AA%E3%83%BC%E3%83%A0%E3%81%AF%E3%80%81PutRecord%E3%81%8A%E3%82%88%E3%81%B3%E3%83%AA%E3%82%AF%E3%82%A8%E3%82%B9%E3%83%88%E3%81%AB%E5%AF%BE%E3%81%97%E3%81%A6%E4%BB%A5%E4%B8%8B%E3%81%AE%E5%90%88%E8%A8%88%E3%82%AF%E3%82%A9%E3%83%BC%E3%82%BF%E3%82%92%E6%8F%90%E4%BE%9B%E3%81%97%E3%81%BE%E3%81%99%E3%80%82PutRecordBatch - if period <= 1*time.Microsecond { - panic("batch interval must be greater than 1 microsecond") +// DrainTimeout bounds delivery after Loop's context ends. +func DrainTimeout(timeout time.Duration) Option { + if timeout <= 0 { + panic("drain timeout must be positive") } + return func(q *Queue) { q.drainTimeout = timeout } +} +// Aggregate concatenates payloads without adding delimiters. Callers supply framing. +// Payloads larger than targetBytes remain separate records, up to the API limit. +func Aggregate(targetBytes int, maxWait time.Duration) Option { + if targetBytes <= 0 || targetBytes > maxRecordBytes || maxWait <= 0 { + panic("invalid aggregation size or wait") + } return func(q *Queue) { - q.batchInterval = period + q.aggregateBytes = targetBytes + q.aggregateWait = maxWait } } -// New return new Queue -func New(fh firehoseiface.FirehoseAPI, DeliveryStreamName string, opts ...Option) *Queue { +// AggregationWaitObserver receives each record's aggregation time before its first attempt. +// The observer must be concurrency-safe and return promptly. +func AggregationWaitObserver(fn func(time.Duration)) Option { + return func(q *Queue) { q.aggregationObserver = fn } +} + +// New creates a queue without starting its delivery workers. +func New(fh firehoseiface.FirehoseAPI, streamName string, opts ...Option) *Queue { q := &Queue{ - queue: make([]*firehose.Record, 0), - workerCh: make(chan []*firehose.Record), firehose: fh, + DeliveryStreamName: aws.String(streamName), initialized: make(chan struct{}), - DeliveryStreamName: aws.String(DeliveryStreamName), + wake: make(chan struct{}, 1), + aggregateWake: make(chan struct{}, 1), + inFlight: make(map[*queuedRecord]struct{}), + para: 1, + maxQueueLength: defaultQueueLength, + maxQueueBytes: defaultQueueBytes, + batchSize: maxBatchRecords, + batchInterval: defaultBatchInterval, + drainTimeout: defaultDrainTimeout, } for _, opt := range opts { opt(q) @@ -90,270 +123,291 @@ func New(fh firehoseiface.FirehoseAPI, DeliveryStreamName string, opts ...Option return q } -// Queue manages a sending list for firehose +// Queue owns an in-memory queue; it does not provide durable or exactly-once delivery. type Queue struct { - queue []*firehose.Record - mu sync.RWMutex - workerCh chan []*firehose.Record - - initialized chan struct{} - initMu sync.Mutex - - para int - inFlightBatchCounter int32 - maxQueueLength int - firehose firehoseiface.FirehoseAPI - errorHandler func(error, []*firehose.Record) - batchSize int - batchInterval time.Duration - DeliveryStreamName *string - - // successCount counts successfully sent records when attemptToSendBatch() is called. - successCount expvar.Int - // retryCount counts retryable records when attemptToSendBatch() is called. - retryCount expvar.Int - // unretryableErrorCount counts non-retryable batch when attemptToSendBatch() is called. - unretryableErrorCount expvar.Int - // queueFullErrorCount counts queue full errors when Enqueue() is called. - queueFullErrorCount expvar.Int - // batchLength is length of the latest batch. - batchLength expvar.Int + mu sync.Mutex + queue []*queuedRecord + current *queuedRecord + inFlight map[*queuedRecord]struct{} + initialized chan struct{} + wake chan struct{} + aggregateWake chan struct{} + started bool + closed bool + para int + maxQueueLength int + maxQueueBytes int + batchSize int + batchInterval time.Duration + drainTimeout time.Duration + aggregateBytes int + aggregateWait time.Duration + aggregationObserver func(time.Duration) + firehose firehoseiface.FirehoseAPI + errorHandler func(error, []*firehose.Record) + stats Stats + + // DeliveryStreamName selects the destination. Do not mutate it after Loop starts. + DeliveryStreamName *string } -// Stats return queue stats -func (q *Queue) Stats() Stats { - return Stats{ - QueueLength: q.len(), - BatchLength: q.batchLength.Value(), - Success: q.successCount.Value(), - RetryCount: q.retryCount.Value(), - UnretryableError: q.unretryableErrorCount.Value(), - QueueFullError: q.queueFullErrorCount.Value(), - } +type queuedRecord struct { + record *firehose.Record + events int64 + createdAt time.Time + sealedAt time.Time + attempted bool } -func (q *Queue) init() error { - q.initMu.Lock() - defer q.initMu.Unlock() +// Enqueue copies the payload. Loop must start within two seconds of the call. +// Shutdown rejects new input; successful acceptance does not imply S3 persistence. +func (q *Queue) Enqueue(record *firehose.Record) error { select { case <-q.initialized: - return fmt.Errorf("already initialized and started the loop") default: + timer := time.NewTimer(startTimeout) + defer timer.Stop() + select { + case <-q.initialized: + case <-timer.C: + return errors.New("loop has not yet started. call Loop before Enqueue") + } } - if q.maxQueueLength == 0 { - q.maxQueueLength = 100000 + q.mu.Lock() + defer q.mu.Unlock() + if q.closed { + return errors.New("queue is closed") } - if q.para == 0 { - q.para = 1 + if record == nil || len(record.Data) == 0 || len(record.Data) > maxRecordBytes { + return errors.New("record data must contain between 1 and 1024000 bytes") } - if q.batchSize == 0 { - q.batchSize = 500 + if q.stats.PendingEvents >= int64(q.maxQueueLength) || int64(len(record.Data)) > int64(q.maxQueueBytes)-q.stats.PendingBytes { + q.stats.QueueFullError++ + return fmt.Errorf("queue capacity exceeded: %d events, %d bytes pending", q.stats.PendingEvents, q.stats.PendingBytes) } - if q.batchInterval == 0 { - q.batchInterval = 2 * time.Millisecond + wasEmpty := len(q.queue) == 0 && q.current == nil + readyWasEmpty := len(q.queue) == 0 + now := time.Now() + q.stats.AcceptedEvents++ + q.stats.PendingEvents++ + q.stats.PendingBytes += int64(len(record.Data)) + if q.aggregateBytes == 0 || len(record.Data) > q.aggregateBytes { + q.sealCurrentLocked(now) + q.queue = append(q.queue, &queuedRecord{ + record: &firehose.Record{Data: append([]byte(nil), record.Data...)}, + events: 1, createdAt: now, sealedAt: now, + }) + } else { + if q.current != nil && (len(q.current.record.Data)+len(record.Data) > q.aggregateBytes || now.Sub(q.current.createdAt) >= q.aggregateWait) { + q.sealCurrentLocked(now) + } + if q.current == nil { + q.current = &queuedRecord{record: &firehose.Record{}, createdAt: now} + select { + case q.aggregateWake <- struct{}{}: + default: + } + } + q.current.record.Data = append(q.current.record.Data, record.Data...) + q.current.events++ + if len(q.current.record.Data) == q.aggregateBytes { + q.sealCurrentLocked(now) + } + } + if wasEmpty || readyWasEmpty && len(q.queue) > 0 { + q.notify() } - close(q.initialized) return nil } -// Loop proceeds jobs in queue -func (q *Queue) Loop(ctx context.Context) error { - if err := q.init(); err != nil { - return err +func (q *Queue) sealCurrentLocked(now time.Time) { + if q.current != nil { + q.current.sealedAt = now + q.queue = append(q.queue, q.current) + q.current = nil + } +} + +func (q *Queue) notify() { + select { + case q.wake <- struct{}{}: + default: } +} +// Loop closes input when ctx ends, then drains within the configured timeout. +// It returns a drain error with remaining counts if delivery cannot finish in time. +func (q *Queue) Loop(ctx context.Context) error { + q.mu.Lock() + if q.started { + q.mu.Unlock() + return errors.New("queue already initialized") + } + q.started = true + close(q.initialized) + q.mu.Unlock() var wg sync.WaitGroup - wg.Add(q.para) + if q.aggregateBytes > 0 { + wg.Add(1) + go func() { + defer wg.Done() + q.aggregateLoop(ctx) + }() + } for i := 0; i < q.para; i++ { + wg.Add(1) go func() { defer wg.Done() q.loop(ctx) }() } + <-ctx.Done() + q.mu.Lock() + q.closed = true + q.sealCurrentLocked(time.Now()) + q.mu.Unlock() + drainCtx, cancel := context.WithTimeout(context.Background(), q.drainTimeout) + defer cancel() wg.Wait() - - // draining process: - var bf = backoff.NewExponentialBackOff() + bf := backoff.NewExponentialBackOff() bf.MaxElapsedTime = 0 - for { - if q.remaining() { - nextInterval := q.sendBatch(bf) - time.Sleep(nextInterval) - continue + if drainCtx.Err() != nil { + stat := q.Stats() + return fmt.Errorf("drain incomplete: %d events, %d bytes pending: %w", stat.PendingEvents, stat.PendingBytes, drainCtx.Err()) } - if !q.remaining() { - break + batch := q.takeBatch() + if len(batch) == 0 { + return nil + } + if q.send(drainCtx, batch) { + timer := time.NewTimer(bf.NextBackOff()) + select { + case <-timer.C: + case <-drainCtx.Done(): + } + timer.Stop() + } else { + bf.Reset() } } - return nil } -func (q *Queue) loop(ctx context.Context) { - var bf = backoff.NewExponentialBackOff() - bf.MaxElapsedTime = 0 - var nextInterval time.Duration = q.batchInterval - ticker := time.NewTicker(q.batchInterval) - defer ticker.Stop() - +func (q *Queue) aggregateLoop(ctx context.Context) { + timer := time.NewTimer(idleInterval) + defer timer.Stop() for { - ticker.Reset(nextInterval) select { case <-ctx.Done(): return - case <-ticker.C: - nextInterval = q.sendBatch(bf) + case <-q.aggregateWake: + case <-timer.C: } - } -} - -func isRetryable(err error) bool { - if reqfailure, ok := err.(awserr.RequestFailure); ok { - // ref. https://github.com/aws/aws-sdk-go/blob/fe72a52350a8962175bb71c531ec9724ce48abd8/aws/request/retryer.go#L228-L250 - switch reqfailure.StatusCode() { - case - 429, - 500, - 502, - 503, - 504: - return true + if !timer.Stop() { + select { + case <-timer.C: + default: + } } + q.mu.Lock() + delay := idleInterval + if q.current != nil { + delay = time.Until(q.current.createdAt.Add(q.aggregateWait)) + if delay <= 0 { + q.sealCurrentLocked(time.Now()) + q.notify() + delay = idleInterval + } + } + q.mu.Unlock() + timer.Reset(delay) } - return request.IsErrorRetryable(err) || - request.IsErrorThrottle(err) || - isErrConnectionResetByPeer(err) -} - -// request.isErrConnectionReset, which is used inside request.IsErrorRetryable, -// intentionally marks false (do not retry) if it contains "read: connection reset". -// ref. https://github.com/aws/aws-sdk-go/blob/7814a7f61bf93cb54d00a5f97918c5501f07d351/aws/request/connection_reset_error.go#L7-L18 -// However, in firehose, the connection reset by peer error often bursts, so it is to be retry -func isErrConnectionResetByPeer(err error) bool { - return strings.Contains(err.Error(), "read: connection reset by peer") } -// Enqueue add data to queue -func (q *Queue) Enqueue(r *firehose.Record) error { - select { - case <-q.initialized: - // nop and continue to send - default: - // It is useless if a timer is created in every Send invocation, so tuck the "default:" here. +func (q *Queue) loop(ctx context.Context) { + bf := backoff.NewExponentialBackOff() + bf.MaxElapsedTime = 0 + timer := time.NewTimer(q.batchInterval) + defer timer.Stop() + var notBefore time.Time + for { select { - case <-time.After(2 * time.Second): - return errors.New("loop has not yet started. call Loop() before Send()") - case <-q.initialized: - // nop and continue to send + case <-ctx.Done(): + return + case <-q.wake: + case <-timer.C: + } + if ctx.Err() != nil { + return + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + if delay := time.Until(notBefore); delay > 0 { + timer.Reset(delay) + continue + } + batch := q.takeBatch() + if len(batch) > 0 && q.send(ctx, batch) { + notBefore = time.Now().Add(bf.NextBackOff()) + timer.Reset(time.Until(notBefore)) + } else { + bf.Reset() + if len(batch) > 0 { + notBefore = time.Now().Add(q.batchInterval) + } + timer.Reset(q.nextDelay()) } } - if l := q.len(); l >= q.maxQueueLength { - q.queueFullErrorCount.Add(1) - return fmt.Errorf("too many jobs accumlated: %d", l) - } - q.push(r) - return nil } -func (q *Queue) sendBatch(bf backoff.BackOff) time.Duration { - if q.len() == 0 { +func (q *Queue) nextDelay() time.Duration { + q.mu.Lock() + defer q.mu.Unlock() + if len(q.queue) > 0 { return q.batchInterval } - batch := q.createBatch() - q.batchLength.Set(int64(len(batch))) - - if err := q.attemptToSendBatch(batch); err != nil { - // Handle non-retryable errors directly. - q.handleError(err, batch) - return bf.NextBackOff() - } - bf.Reset() - return q.batchInterval + return idleInterval } -func (q *Queue) createBatch() []*firehose.Record { +func (q *Queue) takeBatch() []*queuedRecord { q.mu.Lock() - defer q.mu.Unlock() - - batchSize := q.batchSize - if len(q.queue) < batchSize { - batchSize = len(q.queue) + n := 0 + bytes := 0 + for n < len(q.queue) && n < q.batchSize { + if bytes+len(q.queue[n].record.Data) > maxBatchBytes { + break + } + bytes += len(q.queue[n].record.Data) + n++ } - - batch := make([]*firehose.Record, batchSize) - copy(batch, q.queue) - q.queue = q.queue[batchSize:] - return batch -} - -func (q *Queue) attemptToSendBatch(batch []*firehose.Record) error { - atomic.AddInt32(&q.inFlightBatchCounter, 1) - defer atomic.AddInt32(&q.inFlightBatchCounter, -1) - - input := &firehose.PutRecordBatchInput{ - DeliveryStreamName: q.DeliveryStreamName, - Records: batch, + batch := make([]*queuedRecord, n) + copy(batch, q.queue[:n]) + for i := 0; i < n; i++ { + q.queue[i] = nil } - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - resp, err := q.firehose.PutRecordBatchWithContext(ctx, input) - if err != nil { - // Handle non-retryable errors directly. - if !isRetryable(err) { - q.unretryableErrorCount.Add(1) - q.handleError(err, batch) - return nil - } - // Retry failed records. - if resp != nil && resp.FailedPutCount != nil && *resp.FailedPutCount > 0 { - q.successCount.Add(int64(len(batch)) - int64(*resp.FailedPutCount)) - q.retryFailedRecords(batch, resp) - return nil - } - return err + q.queue = q.queue[n:] + if len(q.queue) == 0 { + q.queue = nil + } else { + q.notify() } - - q.successCount.Add(int64(len(batch))) - return nil -} - -func (q *Queue) retryFailedRecords(batch []*firehose.Record, resp *firehose.PutRecordBatchOutput) { - retryRecords := make([]*firehose.Record, 0, *resp.FailedPutCount) - - for i, record := range batch { - if resp.RequestResponses != nil && resp.RequestResponses[i] != nil && resp.RequestResponses[i].ErrorCode != nil { - q.retryCount.Add(1) - retryRecords = append(retryRecords, record) + var waits []time.Duration + for _, record := range batch { + q.inFlight[record] = struct{}{} + if !record.attempted { + if q.aggregateBytes > 0 && q.aggregationObserver != nil { + waits = append(waits, record.sealedAt.Sub(record.createdAt)) + } + record.attempted = true } } - q.push(retryRecords...) -} -func (q *Queue) handleError(err error, r []*firehose.Record) { - if q.errorHandler != nil { - q.errorHandler(err, r) - return + q.mu.Unlock() + for _, wait := range waits { + q.aggregationObserver(wait) } - log.Println(err) -} - -func (q *Queue) len() int { - q.mu.RLock() - defer q.mu.RUnlock() - return len(q.queue) -} - -func (q *Queue) inFlightBatch() bool { - return atomic.LoadInt32(&q.inFlightBatchCounter) > 0 -} - -func (q *Queue) remaining() bool { - // XXX race - return q.len() > 0 || q.inFlightBatch() -} - -func (q *Queue) push(r ...*firehose.Record) { - q.mu.Lock() - defer q.mu.Unlock() - q.queue = append(q.queue, r...) + return batch } diff --git a/firequeue_test.go b/firequeue_test.go index b2b8904..71c2b83 100644 --- a/firequeue_test.go +++ b/firequeue_test.go @@ -1,10 +1,11 @@ package firequeue_test import ( + "bytes" "context" "fmt" - "math/rand" "strings" + "sync" "sync/atomic" "testing" "time" @@ -17,434 +18,565 @@ import ( "github.com/natureglobal/firequeue" ) -type testAWSError struct { +type mockFirehose struct { + firehoseiface.FirehoseAPI + send func(context.Context, *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) } -func (te *testAWSError) Error() string { - return "retryable error" +func (f *mockFirehose) PutRecordBatchWithContext(ctx context.Context, input *firehose.PutRecordBatchInput, _ ...request.Option) (*firehose.PutRecordBatchOutput, error) { + return f.send(ctx, input) } -type errCodeMode string - -const ( - retryable errCodeMode = "retryable" - unretryable errCodeMode = "unretryable" -) - -var ecm = retryable - -func (te *testAWSError) Code() string { - switch ecm { - case retryable: - return request.ErrCodeResponseTimeout - case unretryable: - return request.CanceledErrorCode - default: - return "" +func success(input *firehose.PutRecordBatchInput) *firehose.PutRecordBatchOutput { + entries := make([]*firehose.PutRecordBatchResponseEntry, len(input.Records)) + for i := range entries { + entries[i] = &firehose.PutRecordBatchResponseEntry{RecordId: aws.String("accepted")} } + return &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(0), RequestResponses: entries} } -func (te *testAWSError) Message() string { - return "msg" +type runningQueue struct { + queue *firequeue.Queue + cancel context.CancelFunc + done chan struct{} + err error + allowError bool } -func (te *testAWSError) OrigErr() error { - return fmt.Errorf("retryable error!") +func start(t *testing.T, f *mockFirehose, opts ...firequeue.Option) *runningQueue { + t.Helper() + options := []firequeue.Option{ + firequeue.DrainTimeout(250 * time.Millisecond), + firequeue.ErrorHandler(func(error, []*firehose.Record) {}), + } + options = append(options, opts...) + ctx, cancel := context.WithCancel(context.Background()) + r := &runningQueue{queue: firequeue.New(f, "device-messages", options...), cancel: cancel, done: make(chan struct{})} + go func() { + r.err = r.queue.Loop(ctx) + close(r.done) + }() + t.Cleanup(func() { + if err := r.stop(t); err != nil && !r.allowError { + t.Error(err) + } + }) + return r } -var _ awserr.Error = (*testAWSError)(nil) - -func init() { - rand.Seed(time.Now().UTC().UnixNano()) +func (r *runningQueue) stop(t *testing.T) error { + t.Helper() + r.cancel() + select { + case <-r.done: + return r.err + case <-time.After(2 * time.Second): + t.Fatal("queue did not stop within two seconds") + return nil + } } -type testFirehose struct { - firehoseiface.FirehoseAPI - counter uint32 +func enqueue(t *testing.T, q *firequeue.Queue, data []byte) { + t.Helper() + if err := q.Enqueue(&firehose.Record{Data: data}); err != nil { + t.Fatal(err) + } } -// fhErrorRate is a rate of error occurrence. -// 0 means no error, 10 means always error. -var fhErrorRate = 3 +func waitFor(t *testing.T, condition func() bool) { + t.Helper() + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + timeout := time.NewTimer(2 * time.Second) + defer timeout.Stop() + for !condition() { + select { + case <-ticker.C: + case <-timeout.C: + t.Fatal("condition did not become true") + } + } +} -func (tf *testFirehose) PutRecordBatch(input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { - if rand.Intn(10) >= fhErrorRate { - return &firehose.PutRecordBatchOutput{ - FailedPutCount: aws.Int64(0), - }, nil +func assertBalanced(t *testing.T, s firequeue.Stats) { + t.Helper() + if s.AcceptedEvents != s.SuccessfulEvents+s.FailedEvents+s.PendingEvents { + t.Fatalf("unbalanced events: %+v", s) } +} - var failedPutCounter int64 - inputLength := len(input.Records) - firehoseResp := make([]*firehose.PutRecordBatchResponseEntry, inputLength) - for i := 0; i < inputLength; i++ { - if rand.Intn(10) < fhErrorRate { - firehoseResp[i] = &firehose.PutRecordBatchResponseEntry{ - ErrorCode: aws.String("test_error"), - ErrorMessage: aws.String("test_error_message"), - } - failedPutCounter++ - } else { - firehoseResp[i] = &firehose.PutRecordBatchResponseEntry{ - RecordId: aws.String(fmt.Sprintf("%d", atomic.LoadUint32(&tf.counter))), - } +func TestParallelWorkersWakeForBacklog(t *testing.T) { + entered := make(chan struct{}, 8) + release := make(chan struct{}) + var once sync.Once + t.Cleanup(func() { once.Do(func() { close(release) }) }) + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + entered <- struct{}{} + select { + case <-release: + return success(input), nil + case <-ctx.Done(): + return nil, ctx.Err() } + }} + r := start(t, f, firequeue.Parallel(4), firequeue.BatchSize(1)) + for i := 0; i < 4; i++ { + enqueue(t, r.queue, []byte("event")) } - resp := &firehose.PutRecordBatchOutput{ - RequestResponses: firehoseResp, - FailedPutCount: aws.Int64(failedPutCounter), + for i := 0; i < 4; i++ { + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("not all parallel workers joined the backlog") + } } - - if failedPutCounter > 0 { - return resp, &testAWSError{} + once.Do(func() { close(release) }) + if err := r.stop(t); err != nil { + t.Fatal(err) } - return resp, nil } -func (tf *testFirehose) PutRecordBatchWithContext(_ context.Context, r *firehose.PutRecordBatchInput, _ ...request.Option) (*firehose.PutRecordBatchOutput, error) { - return tf.PutRecordBatch(r) -} - -// success -func TestQueue(t *testing.T) { - fhErrorRate = 3 - testCases := []struct { - name string - times int - }{ - { - name: "1 record", - times: 1, - }, - { - name: "50 records", - times: 50, - }, - { - name: "100 records", - times: 100, - }, - { - name: "1000 records", - times: 1000, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(100*time.Millisecond)) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - go q.Loop(ctx) - - time.Sleep(1000 * time.Millisecond) - - for i := 0; i < tc.times; i++ { - err := q.Enqueue(&firehose.Record{ - Data: []byte("test"), - }) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) - } +func TestAggregationDeadlineWhileDeliveryBlocked(t *testing.T) { + entered := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + t.Cleanup(func() { once.Do(func() { close(release) }) }) + var calls int + var waitsMu sync.Mutex + var waits []time.Duration + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + if calls == 1 { + close(entered) + select { + case <-release: + case <-ctx.Done(): + return nil, ctx.Err() } - time.Sleep(10 * time.Second) - stats := q.Stats() - if stats.Success != int64(tc.times) { - t.Log(stats) - t.Errorf("Not all records were success. expected: %d, actual: %d", tc.times, stats.Success) - } - t.Log(q.Stats()) - - }) + } + return success(input), nil + }} + r := start(t, f, firequeue.Aggregate(16, 20*time.Millisecond), firequeue.AggregationWaitObserver(func(d time.Duration) { + waitsMu.Lock() + waits = append(waits, d) + waitsMu.Unlock() + })) + enqueue(t, r.queue, bytes.Repeat([]byte("x"), 16)) + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("first request did not start") } -} - -func TestQueue_Para(t *testing.T) { - testCases := []struct { - name string - times int - para int - }{ - { - name: "parallel 2", - times: 1000, - para: 2, - }, - { - name: "parallel 10", - times: 1000, - para: 10, - }, - { - name: "parallel 100", - times: 1000, - para: 100, - }, + enqueue(t, r.queue, []byte("y")) + time.Sleep(150 * time.Millisecond) + once.Do(func() { close(release) }) + waitFor(t, func() bool { return r.queue.Stats().SuccessfulEvents == 2 }) + if err := r.stop(t); err != nil { + t.Fatal(err) } + waitsMu.Lock() + defer waitsMu.Unlock() + if len(waits) != 2 || waits[1] >= 100*time.Millisecond { + t.Fatalf("aggregation waited for blocked delivery: %v", waits) + } +} - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(1*time.Second), firequeue.Parallel(tc.para)) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - go q.Loop(ctx) - - time.Sleep(1000 * time.Millisecond) - - for i := 0; i < tc.times; i++ { - err := q.Enqueue(&firehose.Record{ - Data: []byte("test"), - }) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) +func TestAggregationPreservesPayloads(t *testing.T) { + for _, enabled := range []bool{false, true} { + t.Run(fmt.Sprintf("enabled=%t", enabled), func(t *testing.T) { + var mu sync.Mutex + var got [][]byte + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + mu.Lock() + defer mu.Unlock() + for _, record := range input.Records { + got = append(got, append([]byte(nil), record.Data...)) } + return success(input), nil + }} + var opts []firequeue.Option + if enabled { + opts = append(opts, firequeue.Aggregate(4, time.Hour)) } - time.Sleep(10 * time.Second) - stats := q.Stats() - if stats.Success != int64(tc.times) { - t.Log(stats) - t.Errorf("Not all records were success. expected: %d, actual: %d", tc.times, stats.Success) + r := start(t, f, opts...) + for _, data := range []string{"a\n", "b\n", "c\n"} { + enqueue(t, r.queue, []byte(data)) + } + if err := r.stop(t); err != nil { + t.Fatal(err) + } + mu.Lock() + defer mu.Unlock() + wantRecords := 3 + if enabled { + wantRecords = 2 + } + if len(got) != wantRecords || string(bytes.Join(got, nil)) != "a\nb\nc\n" { + t.Fatalf("unexpected records: %q", got) + } + s := r.queue.Stats() + assertBalanced(t, s) + if s.SuccessfulEvents != 3 || s.Success != int64(wantRecords) || s.PendingBytes != 0 { + t.Fatalf("unexpected stats: %+v", s) } - t.Log(q.Stats()) - }) } } -func TestQueue_CallLoopTwoTimes(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env") - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - go q.Loop(ctx) - - time.Sleep(1000 * time.Millisecond) - err := q.Loop(ctx) - if err == nil || !strings.Contains(err.Error(), "already initialized") { - fmt.Println(err) - t.Errorf("already initialized error should be occurred but: %s", err) +func TestAggregationBoundariesAndOwnership(t *testing.T) { + var got [][]byte + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + for _, record := range input.Records { + got = append(got, append([]byte(nil), record.Data...)) + } + return success(input), nil + }} + r := start(t, f, firequeue.Aggregate(65536, time.Hour)) + var want [][]byte + for _, size := range []int{65535, 65536, 65537, 1024000} { + payload := bytes.Repeat([]byte("x"), size) + want = append(want, append([]byte(nil), payload...)) + enqueue(t, r.queue, payload) + payload[0] = 'z' } -} - -func TestQueue_SendWithoutLoop(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env") - - err := q.Enqueue(nil) - if err == nil || !strings.Contains(err.Error(), "loop has not yet started") { - t.Errorf("loop has not yet started error should be occurred but: %s", err) + for _, record := range []*firehose.Record{nil, {}, {Data: make([]byte, 1024001)}} { + if err := r.queue.Enqueue(record); err == nil { + t.Fatal("invalid input was accepted") + } + } + if err := r.stop(t); err != nil { + t.Fatal(err) + } + if len(got) != len(want) { + t.Fatalf("got %d records, want %d", len(got), len(want)) + } + for i := range want { + if !bytes.Equal(got[i], want[i]) { + t.Fatalf("record %d changed", i) + } } } -func TestQueue_ExceedBatchSize(t *testing.T) { - defer func() { - if r := recover(); r == nil { - t.Errorf("panic should be occurred but not") +func TestAggregationDeadlineDoesNotSlide(t *testing.T) { + delivered := make(chan struct{}, 1) + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + select { + case delivered <- struct{}{}: + default: + } + return success(input), nil + }} + r := start(t, f, firequeue.Aggregate(65536, 30*time.Millisecond)) + enqueue(t, r.queue, []byte("first\n")) + producerCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + producerDone := make(chan struct{}) + go func() { + defer close(producerDone) + ticker := time.NewTicker(5 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-producerCtx.Done(): + return + case <-ticker.C: + if err := r.queue.Enqueue(&firehose.Record{Data: []byte("next\n")}); err != nil { + return + } + } } }() - tf := &testFirehose{} - firequeue.New(tf, "env", firequeue.BatchSize(501)) + select { + case <-delivered: + case <-time.After(150 * time.Millisecond): + t.Error("continuous input postponed the first flush") + } + cancel() + <-producerDone } -func TestQueue_QueueLength(t *testing.T) { - defer func() { - if r := recover(); r == nil { - t.Errorf("panic should be occurred but not") +func TestHTTPPartialFailureRetriesOnlyFailedRecord(t *testing.T) { + entered := make(chan struct{}) + release := make(chan struct{}) + var calls int + var got []string + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + for _, record := range input.Records { + got = append(got, string(record.Data)) } - }() - tf := &testFirehose{} - firequeue.New(tf, "env", firequeue.MaxQueueLength(-1)) + if calls == 1 { + close(entered) + select { + case <-release: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + resp := success(input) + if calls == 2 { + if len(input.Records) != 2 { + return nil, fmt.Errorf("expected two records, got %d", len(input.Records)) + } + resp.FailedPutCount = aws.Int64(1) + resp.RequestResponses[1] = &firehose.PutRecordBatchResponseEntry{ErrorCode: aws.String("InternalFailure")} + } + return resp, nil + }} + r := start(t, f, firequeue.Aggregate(4, time.Hour)) + enqueue(t, r.queue, []byte("warm")) + <-entered + for _, s := range []string{"a\n", "b\n", "c\n", "d\n"} { + enqueue(t, r.queue, []byte(s)) + } + close(release) + waitFor(t, func() bool { return r.queue.Stats().RetryCount == 1 }) + if err := r.stop(t); err != nil { + t.Fatal(err) + } + if strings.Join(got, "|") != "warm|a\nb\n|c\nd\n|c\nd\n" { + t.Fatalf("unexpected retry: %q", got) + } + s := r.queue.Stats() + assertBalanced(t, s) + if s.SuccessfulEvents != 5 || s.Success != 3 || s.SentRecords != 4 || s.EstimatedBillableBytes != 4*5120 { + t.Fatalf("unexpected stats: %+v", s) + } } -func TestQueue_DrainProcess(t *testing.T) { - tests := []struct { - name string - count int +func TestBatchLimits(t *testing.T) { + for _, tc := range []struct { + name string + size, count, largest int }{ - {"10 records", 10}, - {"50 records", 50}, - {"100 records", 100}, - {"1000 records", 1000}, - {"2000 records", 2000}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(1*time.Second)) - ctx, cancel := context.WithCancel(context.Background()) - go q.Loop(ctx) - - time.Sleep(1000 * time.Millisecond) - - for i := 0; i < tt.count; i++ { - err := q.Enqueue(&firehose.Record{ - Data: []byte("test"), - }) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) + {"record count", 1, 502, 500}, + {"payload bytes", 65536, 64, 61}, + } { + t.Run(tc.name, func(t *testing.T) { + entered, release := make(chan struct{}), make(chan struct{}) + calls, largest := 0, 0 + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + if calls == 1 { + close(entered) + select { + case <-release: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if len(input.Records) > largest { + largest = len(input.Records) + } + var size int + for _, record := range input.Records { + size += len(record.Data) } + if size > 4000000 || len(input.Records) > 500 { + t.Errorf("invalid request: %d records, %d bytes", len(input.Records), size) + } + return success(input), nil + }} + r := start(t, f) + enqueue(t, r.queue, bytes.Repeat([]byte("x"), tc.size)) + <-entered + for i := 1; i < tc.count; i++ { + enqueue(t, r.queue, bytes.Repeat([]byte("x"), tc.size)) + } + close(release) + if err := r.stop(t); err != nil { + t.Fatal(err) } - cancel() - time.Sleep(10 * time.Second) - stats := q.Stats() - if stats.Success != int64(tt.count) { - t.Log(stats) - t.Errorf("Not all records were success. expected: %d, actual: %d", tt.count, stats.Success) + if largest != tc.largest || r.queue.Stats().SuccessfulEvents != int64(tc.count) { + t.Fatalf("largest batch %d, stats %+v", largest, r.queue.Stats()) } - t.Log(q.Stats()) }) } } -func TestQueue_DrainProcessPara(t *testing.T) { - tests := []struct { - name string - count int - para int +func TestCapacityIncludesInFlightAndConcurrentInput(t *testing.T) { + for _, tc := range []struct { + name string + events, bytes, capacity int }{ - {"parallel 2", 1000, 2}, - {"parallel 10", 1000, 10}, - {"parallel 100", 1000, 100}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(1*time.Second)) - ctx, cancel := context.WithCancel(context.Background()) - go q.Loop(ctx) - - time.Sleep(1000 * time.Millisecond) - - for i := 0; i < tt.count; i++ { - err := q.Enqueue(&firehose.Record{ - Data: []byte("test"), - }) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) + {"event limit", 5, 1000, 5}, + {"byte limit", 1000, 8, 4}, + } { + t.Run(tc.name, func(t *testing.T) { + entered, release := make(chan struct{}), make(chan struct{}) + var once sync.Once + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + once.Do(func() { close(entered) }) + select { + case <-release: + return success(input), nil + case <-ctx.Done(): + return nil, ctx.Err() } + }} + r := start(t, f, firequeue.MaxQueueLength(tc.events), firequeue.MaxQueueBytes(tc.bytes), firequeue.Parallel(2)) + enqueue(t, r.queue, []byte("x\n")) + <-entered + var wg sync.WaitGroup + var accepted int64 + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if err := r.queue.Enqueue(&firehose.Record{Data: []byte("y\n")}); err == nil { + atomic.AddInt64(&accepted, 1) + } + }() } - cancel() - time.Sleep(10 * time.Second) - stats := q.Stats() - if stats.Success != int64(tt.count) { - t.Log(stats) - t.Errorf("Not all records were success. expected: %d, actual: %d", tt.count, stats.Success) + wg.Wait() + s := r.queue.Stats() + assertBalanced(t, s) + if accepted+1 != int64(tc.capacity) || s.PendingBytes != int64(2*tc.capacity) || s.QueueFullError != int64(21-tc.capacity) { + t.Fatalf("capacity escaped: %+v, accepted %d", s, accepted) + } + close(release) + if err := r.stop(t); err != nil { + t.Fatal(err) + } + if err := r.queue.Enqueue(&firehose.Record{Data: []byte("late")}); err == nil { + t.Fatal("accepted input after shutdown") } - t.Log(q.Stats()) }) } } -// Test for stats work correctly. -var qLength bool - -func TestStats(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(1*time.Second)) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - go q.Loop(ctx) - time.Sleep(1000 * time.Millisecond) - - for i := 0; i < 100; i++ { - err := q.Enqueue(&firehose.Record{ - Data: []byte("test"), - }) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) +func TestCancellationRetriesAmbiguousRequestDuringDrain(t *testing.T) { + entered := make(chan struct{}) + var calls int + f := &mockFirehose{send: func(ctx context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + if calls == 1 { + close(entered) + <-ctx.Done() + return nil, ctx.Err() } - if q.Stats().QueueLength > 0 { - qLength = true + if ctx.Err() != nil { + t.Error("drain used the canceled producer context") } + return success(input), nil + }} + r := start(t, f) + enqueue(t, r.queue, []byte("payload\n")) + <-entered + if err := r.stop(t); err != nil { + t.Fatal(err) } - time.Sleep(10 * time.Second) - stats := q.Stats() - if stats.Success != 100 { - t.Log(stats) - t.Errorf("Not all records were success. expected: %d, actual: %d", 100, stats.Success) + s := r.queue.Stats() + assertBalanced(t, s) + if calls != 2 || s.SuccessfulEvents != 1 || s.RetryCount != 1 { + t.Fatalf("unexpected cancellation accounting: %+v", s) } - if stats.BatchLength == 0 { - t.Errorf("BatchLength should be more than 0 but: %d", stats.BatchLength) - } - if !qLength { - t.Errorf("QueueLength should be true: %t", qLength) - } - t.Log(q.Stats()) } -func TestStats_QueueFullError(t *testing.T) { - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchInterval(10*time.Second), firequeue.MaxQueueLength(1)) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - go q.Loop(ctx) - time.Sleep(1000 * time.Millisecond) - - err := q.Enqueue(&firehose.Record{}) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) +func TestDrainTimeoutRetainsUnacknowledgedEvents(t *testing.T) { + entered := make(chan struct{}) + var once sync.Once + f := &mockFirehose{send: func(ctx context.Context, _ *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + once.Do(func() { close(entered) }) + <-ctx.Done() + return nil, ctx.Err() + }} + r := start(t, f, firequeue.DrainTimeout(30*time.Millisecond)) + r.allowError = true + enqueue(t, r.queue, []byte("retained\n")) + <-entered + err := r.stop(t) + if err == nil || !strings.Contains(err.Error(), "1 events, 9 bytes pending") { + t.Fatalf("unexpected drain result: %v", err) } - err = q.Enqueue(&firehose.Record{}) - if err == nil { - t.Errorf("error should be occurred but not") + s := r.queue.Stats() + assertBalanced(t, s) + if s.PendingEvents != 1 || s.SuccessfulEvents != 0 || s.FailedEvents != 0 { + t.Fatalf("unacknowledged event disappeared: %+v", s) } - stats := q.Stats() - if stats.QueueFullError != 1 { - t.Errorf("QueueFullError should be 1 but: %d", stats.QueueFullError) - } - t.Log(stats) } -func TestStat_RetryCount(t *testing.T) { - // It means always return retryable error. - fhErrorRate = 10 - tf := &testFirehose{} - q := firequeue.New(tf, "env") - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - go q.Loop(ctx) - time.Sleep(1000 * time.Millisecond) - for i := 0; i < 10; i++ { - err := q.Enqueue(&firehose.Record{}) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) - } - } - time.Sleep(5 * time.Second) - stats := q.Stats() - if stats.RetryCount == 0 { - t.Errorf("retryCount should be more than 0 but: %d", stats.RetryCount) +func TestMalformedResponsesDoNotAcknowledgeEvents(t *testing.T) { + for _, tc := range []struct { + name string + response *firehose.PutRecordBatchOutput + }{ + {"nil response", nil}, + {"missing count", &firehose.PutRecordBatchOutput{}}, + {"missing entry", &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(0)}}, + {"nil entry", &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(0), RequestResponses: []*firehose.PutRecordBatchResponseEntry{nil}}}, + {"missing ID", &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(0), RequestResponses: []*firehose.PutRecordBatchResponseEntry{{}}}}, + {"empty ID", &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(0), RequestResponses: []*firehose.PutRecordBatchResponseEntry{{RecordId: aws.String("")}}}}, + {"inconsistent count", &firehose.PutRecordBatchOutput{FailedPutCount: aws.Int64(1), RequestResponses: []*firehose.PutRecordBatchResponseEntry{{}}}}, + } { + t.Run(tc.name, func(t *testing.T) { + var calls int + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + if calls == 1 { + return tc.response, nil + } + return success(input), nil + }} + r := start(t, f) + enqueue(t, r.queue, []byte("event\n")) + waitFor(t, func() bool { return r.queue.Stats().RetryCount == 1 }) + if r.queue.Stats().SuccessfulEvents != 0 { + t.Fatal("malformed response acknowledged an event") + } + if err := r.stop(t); err != nil { + t.Fatal(err) + } + if r.queue.Stats().SuccessfulEvents != 1 { + t.Fatal("retry did not recover") + } + }) } - t.Log(stats) } -func TestStat_UnRetryableErrorCount(t *testing.T) { - // It means always return unretryable error. - fhErrorRate = 10 - ecm = unretryable - - tf := &testFirehose{} - q := firequeue.New(tf, "env", firequeue.BatchSize(1)) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() +func TestRetryableAndTerminalErrors(t *testing.T) { + for _, tc := range []struct { + name string + err error + retry bool + }{ + {"429", awserr.NewRequestFailure(awserr.New("throttled", "", nil), 429, ""), true}, + {"500", awserr.NewRequestFailure(awserr.New("failed", "", nil), 500, ""), true}, + {"timeout", awserr.New(request.ErrCodeResponseTimeout, "timeout", nil), true}, + {"connection reset", fmt.Errorf("read: connection reset by peer"), true}, + {"invalid argument", awserr.New("InvalidArgumentException", "invalid", nil), false}, + } { + t.Run(tc.name, func(t *testing.T) { + calls := 0 + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + calls++ + if calls == 1 { + return nil, tc.err + } + return success(input), nil + }} + r := start(t, f) + enqueue(t, r.queue, []byte("event\n")) + waitFor(t, func() bool { s := r.queue.Stats(); return s.RetryCount > 0 || s.FailedEvents > 0 }) + if err := r.stop(t); err != nil { + t.Fatal(err) + } + s := r.queue.Stats() + assertBalanced(t, s) + if tc.retry && (s.SuccessfulEvents != 1 || calls != 2) || !tc.retry && (s.FailedEvents != 1 || calls != 1) { + t.Fatalf("unexpected disposition: %+v, calls %d", s, calls) + } + }) + } +} - go q.Loop(ctx) - time.Sleep(1000 * time.Millisecond) - err := q.Enqueue(&firehose.Record{}) - if err != nil { - t.Errorf("error should not be occurred but: %s", err) - } - time.Sleep(5 * time.Second) - stats := q.Stats() - if stats.UnretryableError == 0 { - t.Errorf("retryCount should be more than 0 but: %d", stats.RetryCount) - } - t.Log(stats) +func TestLoopCannotRestart(t *testing.T) { + f := &mockFirehose{send: func(_ context.Context, input *firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { + return success(input), nil + }} + r := start(t, f) + enqueue(t, r.queue, []byte("event\n")) + if err := r.queue.Loop(context.Background()); err == nil || !strings.Contains(err.Error(), "already initialized") { + t.Fatalf("second Loop: %v", err) + } } diff --git a/stats.go b/stats.go index 92e9f70..257f1b5 100644 --- a/stats.go +++ b/stats.go @@ -1,11 +1,64 @@ package firequeue -// Stats contains firequeue statistics. +import "time" + +// Stats separates input events from Firehose records and reports pending payloads. type Stats struct { - QueueLength int - BatchLength int64 - Success int64 - RetryCount int64 + // QueueLength counts pending events outside active requests. + QueueLength int + // BatchLength counts Firehose records in the most recent request. + BatchLength int64 + // Success counts acknowledged Firehose records, not input events. + Success int64 + // RetryCount counts records queued for another attempt. + RetryCount int64 + // UnretryableError counts terminally failed Firehose records. UnretryableError int64 - QueueFullError int64 + // QueueFullError counts input rejections caused by either capacity limit. + QueueFullError int64 + // AcceptedEvents counts input events accepted by Enqueue. + AcceptedEvents int64 + // SuccessfulEvents counts input events acknowledged by Firehose. + SuccessfulEvents int64 + // FailedEvents counts input events with terminal API failures. + FailedEvents int64 + // PendingEvents includes aggregation, queued, retrying, and in-flight events. + PendingEvents int64 + // PendingBytes counts payload bytes owned by all pending events. + PendingBytes int64 + // SentRecords counts every attempted record, including retries. + SentRecords int64 + // EstimatedBillableBytes rounds every attempt up to 5 KiB; it is not a bill. + EstimatedBillableBytes int64 + // OldestEventAge measures the age of the oldest pending input event. + OldestEventAge time.Duration +} + +// Stats returns a concurrency-safe snapshot of delivery accounting. +func (q *Queue) Stats() Stats { + q.mu.Lock() + defer q.mu.Unlock() + stat := q.stats + var inFlightEvents int64 + var oldest time.Time + visit := func(record *queuedRecord) { + if oldest.IsZero() || record.createdAt.Before(oldest) { + oldest = record.createdAt + } + } + for record := range q.inFlight { + inFlightEvents += record.events + visit(record) + } + for _, record := range q.queue { + visit(record) + } + if q.current != nil { + visit(q.current) + } + stat.QueueLength = int(stat.PendingEvents - inFlightEvents) + if !oldest.IsZero() { + stat.OldestEventAge = time.Since(oldest) + } + return stat }