Repository: beatlabs/patron
Pull request #227 (2020-06-25 12:31:13)

Consume Kafka messages since a given duration
## Which problem is this PR solving?

https://github.com/beatlabs/patron/issues/226

## Short description of the changes

* New option `WithDurationOffset`
* New feature to retrieve the offsets per partition given a duration
* Change in simple `Consume()` to consume from the offsets previously retrieved

codecov[bot]: # [Codecov](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=h1) Report
> Merging [#227](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=desc) into [master](https://codecov.io/gh/beatlabs/patron/commit/375084b0e6deed91f444975c248bcbefaad2c2e0&el=desc) will **decrease** coverage by `11.65%`.
> The diff coverage is `48.92%`.

[![Impacted file tree graph](https://codecov.io/gh/beatlabs/patron/pull/227/graphs/tree.svg?width=650&height=150&src=pr&token=sxY15rXW1X)](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           master     #227       +/-   ##
===========================================
- Coverage   84.48%   72.83%   -11.66%     
===========================================
  Files          52       58        +6     
  Lines        3074     3368      +294     
===========================================
- Hits         2597     2453      -144     
- Misses        398      850      +452     
+ Partials       79       65       -14     
```


| [Impacted Files](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [client/amqp/amqp.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2FtcXAvYW1xcC5nbw==) | `40.98% <ø> (ø)` | |
| [client/amqp/option.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2FtcXAvb3B0aW9uLmdv) | `66.66% <ø> (ø)` | |
| [client/amqp/user.pb.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2FtcXAvdXNlci5wYi5nbw==) | `26.92% <ø> (ø)` | |
| [client/http/option.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2h0dHAvb3B0aW9uLmdv) | `100.00% <ø> (ø)` | |
| [client/kafka/async\_producer.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2thZmthL2FzeW5jX3Byb2R1Y2VyLmdv) | `0.00% <0.00%> (ø)` | |
| [client/kafka/sync\_producer.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L2thZmthL3N5bmNfcHJvZHVjZXIuZ28=) | `0.00% <0.00%> (ø)` | |
| [client/redis/redis.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L3JlZGlzL3JlZGlzLmdv) | `44.44% <ø> (ø)` | |
| [client/sns/message.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L3Nucy9tZXNzYWdlLmdv) | `97.53% <ø> (ø)` | |
| [client/sns/publisher.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L3Nucy9wdWJsaXNoZXIuZ28=) | `93.33% <ø> (ø)` | |
| [client/sql/sql.go](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree#diff-Y2xpZW50L3NxbC9zcWwuZ28=) | `17.53% <ø> (-68.73%)` | :arrow_down: |
| ... and [75 more](https://codecov.io/gh/beatlabs/patron/pull/227/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=footer). Last update [d51c80b...9b336eb](https://codecov.io/gh/beatlabs/patron/pull/227?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

mantzas: @teivah CI is failing
teivah: > @teivah CI is failing

@mantzas A good example of https://github.com/beatlabs/patron/issues/228

```
Not equal: 
  expected: []string{"2020-06-29T10:10:14Z", "2020-06-29T11:10:14Z", "2020-06-29T12:10:14Z"}
  actual  : []string{"2020-06-29T12:10:14Z", "2020-06-29T10:10:14Z", "2020-06-29T11:10:14Z"}
```

mantzas: > > @teivah CI is failing
> 
> @mantzas A good example of #228
> 
> ```
> Not equal: 
>   expected: []string{"2020-06-29T10:10:14Z", "2020-06-29T11:10:14Z", "2020-06-29T12:10:14Z"}
>   actual  : []string{"2020-06-29T12:10:14Z", "2020-06-29T10:10:14Z", "2020-06-29T11:10:14Z"}
> ```

@teivah so we should probably merge #228 first correct. Agree?
teivah: > @teivah so we should probably merge #228 first correct. Agree?

@mantzas Yes, I think so.
teivah: @mantzas It did fix the test, so we were correct about it 🎉 

Patch:
diff --git a/component/async/kafka/kafka.go b/component/async/kafka/kafka.go
index eba362ee7..ca615f472 100644
--- a/component/async/kafka/kafka.go
+++ b/component/async/kafka/kafka.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/beatlabs/patron/component/async"
@@ -32,6 +33,9 @@ var topicPartitionOffsetDiff *prometheus.GaugeVec
 var messageStatus *prometheus.CounterVec
 var messageConfirmation *prometheus.CounterVec
 
+// TimeExtractor defines a function extracting a time from a Kafka message.
+type TimeExtractor func(*sarama.ConsumerMessage) (time.Time, error)
+
 // TopicPartitionOffsetDiffGaugeSet creates a new Gauge that measures partition offsets.
 func TopicPartitionOffsetDiffGaugeSet(group, topic string, partition int32, high, offset int64) {
 	topicPartitionOffsetDiff.WithLabelValues(group, topic, strconv.FormatInt(int64(partition), 10)).Set(float64(high - offset))
@@ -84,10 +88,13 @@ func init() {
 
 // ConsumerConfig is the common configuration of patron kafka consumers.
 type ConsumerConfig struct {
-	Brokers      []string
-	Buffer       int
-	DecoderFunc  encoding.DecodeRawFunc
-	SaramaConfig *sarama.Config
+	Brokers               []string
+	Buffer                int
+	DecoderFunc           encoding.DecodeRawFunc
+	DurationBasedConsumer bool
+	DurationOffset        time.Duration
+	TimeExtractor         func(*sarama.ConsumerMessage) (time.Time, error)
+	SaramaConfig          *sarama.Config
 }
 
 type message struct {
diff --git a/component/async/kafka/option.go b/component/async/kafka/option.go
index 85feb791e..9e4d779d6 100644
--- a/component/async/kafka/option.go
+++ b/component/async/kafka/option.go
@@ -90,3 +90,20 @@ func DecoderJSON() OptionFunc {
 		return nil
 	}
 }
+
+// WithDurationOffset allows creating a consumer from a given duration.
+// It accepts a function indicating how to extract the time from a Kafka message.
+func WithDurationOffset(since time.Duration, timeExtractor TimeExtractor) OptionFunc {
+	return func(c *ConsumerConfig) error {
+		if since < 0 {
+			return errors.New("duration must be positive")
+		}
+		if timeExtractor == nil {
+			return errors.New("empty time extractor function")
+		}
+		c.DurationBasedConsumer = true
+		c.DurationOffset = since
+		c.TimeExtractor = timeExtractor
+		return nil
+	}
+}
diff --git a/component/async/kafka/simple/duration_client.go b/component/async/kafka/simple/duration_client.go
new file mode 100644
index 000000000..0b1f77aca
--- /dev/null
+++ b/component/async/kafka/simple/duration_client.go
@@ -0,0 +1,130 @@
+package simple
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/beatlabs/patron/component/async/kafka"
+	"github.com/beatlabs/patron/log"
+)
+
+type durationClient struct {
+	client durationKafkaClientAPI
+}
+
+func newDurationClient(client durationKafkaClientAPI) (durationClient, error) {
+	if client == nil {
+		return durationClient{}, errors.New("empty client api")
+	}
+	return durationClient{client: client}, nil
+}
+
+func (d durationClient) getTimeBasedOffsetsPerPartition(ctx context.Context, topic string, since time.Time, timeExtractor kafka.TimeExtractor) (map[int32]int64, error) {
+	partitionIDs, err := d.client.getPartitionIDs(topic)
+	if err != nil {
+		return nil, err
+	}
+
+	responseCh := make(chan partitionOffsetResponse, len(partitionIDs))
+	d.triggerWorkers(ctx, topic, since, timeExtractor, partitionIDs, responseCh)
+	return d.aggregateResponses(ctx, partitionIDs, responseCh)
+}
+
+type partitionOffsetResponse struct {
+	partitionID int32
+	offset      int64
+	err         error
+}
+
+func (d durationClient) triggerWorkers(ctx context.Context, topic string, since time.Time, timeExtractor kafka.TimeExtractor, partitionIDs []int32, responseCh chan<- partitionOffsetResponse) {
+	for _, partitionID := range partitionIDs {
+		partitionID := partitionID
+		go func() {
+			offset, err := d.getTimeBasedOffset(ctx, topic, since, partitionID, timeExtractor)
+			select {
+			case <-ctx.Done():
+				return
+			case responseCh <- partitionOffsetResponse{
+				partitionID: partitionID,
+				offset:      offset,
+				err:         err,
+			}:
+			}
+		}()
+	}
+}
+
+func (d durationClient) aggregateResponses(ctx context.Context, partitionIDs []int32, responseCh <-chan partitionOffsetResponse) (map[int32]int64, error) {
+	numberOfPartitions := len(partitionIDs)
+	offsets := make(map[int32]int64, numberOfPartitions)
+	numberOfResponses := 0
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, fmt.Errorf("context cancelled before collecting partition responses: %w", ctx.Err())
+		case response := <-responseCh:
+			if response.err != nil {
+				return nil, response.err
+			}
+
+			offsets[response.partitionID] = response.offset
+			numberOfResponses++
+			if numberOfResponses == numberOfPartitions {
+				return offsets, nil
+			}
+		}
+	}
+}
+
+func (d durationClient) getTimeBasedOffset(ctx context.Context, topic string, since time.Time, partitionID int32, timeExtractor kafka.TimeExtractor) (int64, error) {
+	left, err := d.client.getOldestOffset(topic, partitionID)
+	if err != nil {
+		return 0, err
+	}
+
+	newestOffset, err := d.client.getNewestOffset(topic, partitionID)
+	if err != nil {
+		return 0, err
+	}
+	// The right boundary must be inclusive
+	right := newestOffset - 1
+
+	return d.offsetBinarySearch(ctx, topic, since, partitionID, timeExtractor, left, right)
+}
+
+func (d durationClient) offsetBinarySearch(ctx context.Context, topic string, since time.Time, partitionID int32, timeExtractor kafka.TimeExtractor, left, right int64) (int64, error) {
+	for left <= right {
+		mid := left + (right-left)/2
+
+		msg, err := d.client.getMessageAtOffset(ctx, topic, partitionID, mid)
+		if err != nil {
+			// Under extraordinary circumstances (e.g. the retention policy being applied just before retrieving the message at a particular offset),
+			// the offset might not be accessible anymore.
+			// In this case, we simply log a warning and restrict the interval to the right.
+			if errors.Is(err, &outOfRangeOffsetError{}) {
+				log.Warnf("offset %d on partition %d is out of range: %v", mid, partitionID, err)
+				left = mid + 1
+				continue
+			}
+			return 0, fmt.Errorf("error while retrieving message offset %d on partition %d: %w", mid, partitionID, err)
+		}
+
+		t, err := timeExtractor(msg)
+		if err != nil {
+			return 0, fmt.Errorf("error while executing comparator: %w", err)
+		}
+
+		if t.Equal(since) {
+			return mid, nil
+		}
+		if t.Before(since) {
+			left = mid + 1
+		} else {
+			right = mid - 1
+		}
+	}
+
+	return left, nil
+}
diff --git a/component/async/kafka/simple/duration_kafka.go b/component/async/kafka/simple/duration_kafka.go
new file mode 100644
index 000000000..1404f76ad
--- /dev/null
+++ b/component/async/kafka/simple/duration_kafka.go
@@ -0,0 +1,105 @@
+package simple
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/Shopify/sarama"
+)
+
+type outOfRangeOffsetError struct {
+	message string
+}
+
+func (e *outOfRangeOffsetError) Error() string {
+	return e.message
+}
+
+func (e *outOfRangeOffsetError) Is(target error) bool {
+	_, ok := target.(*outOfRangeOffsetError)
+	return ok
+}
+
+type durationKafkaClientAPI interface {
+	getPartitionIDs(topic string) ([]int32, error)
+	getOldestOffset(topic string, partitionID int32) (int64, error)
+	getNewestOffset(topic string, partitionID int32) (int64, error)
+	getMessageAtOffset(ctx context.Context, topic string, partitionID int32, offset int64) (*sarama.ConsumerMessage, error)
+}
+
+type durationKafkaClient struct {
+	kafkaClient     sarama.Client
+	kafkaConsumer   sarama.Consumer
+	consumerTimeout time.Duration
+}
+
+func newDurationKafkaClient(kafkaClient sarama.Client, kafkaConsumer sarama.Consumer, consumerTimeout time.Duration) (durationKafkaClient, error) {
+	if kafkaClient == nil {
+		return durationKafkaClient{}, errors.New("kafka client is nil")
+	}
+	if kafkaConsumer == nil {
+		return durationKafkaClient{}, errors.New("kafka consumer is nil")
+	}
+
+	return durationKafkaClient{
+		kafkaClient:     kafkaClient,
+		kafkaConsumer:   kafkaConsumer,
+		consumerTimeout: consumerTimeout,
+	}, nil
+}
+
+func (c durationKafkaClient) getPartitionIDs(topic string) ([]int32, error) {
+	partitionIDs, err := c.kafkaConsumer.Partitions(topic)
+	if err != nil {
+		return nil, fmt.Errorf("unable to query Kafka to retrieve the partitions of the topic %s: %w", topic, err)
+	}
+
+	return partitionIDs, nil
+}
+
+func (c durationKafkaClient) getOldestOffset(topic string, partitionID int32) (int64, error) {
+	offset, err := c.kafkaClient.GetOffset(topic, partitionID, sarama.OffsetOldest)
+	if err != nil {
+		return 0, fmt.Errorf("error while retrieving oldest offset of partition %d: %w", partitionID, err)
+	}
+	return offset, nil
+}
+
+func (c durationKafkaClient) getNewestOffset(topic string, partitionID int32) (int64, error) {
+	offset, err := c.kafkaClient.GetOffset(topic, partitionID, sarama.OffsetNewest)
+	if err != nil {
+		return 0, fmt.Errorf("error while retrieving newest offset of partition %d: %w", partitionID, err)
+	}
+	return offset, nil
+}
+
+func (c durationKafkaClient) getMessageAtOffset(ctx context.Context, topic string, partitionID int32, offset int64) (*sarama.ConsumerMessage, error) {
+	pc, err := c.kafkaConsumer.ConsumePartition(topic, partitionID, offset)
+	if err != nil {
+		if err == sarama.ErrOffsetOutOfRange {
+			closePartitionConsumer(pc)
+			return nil, &outOfRangeOffsetError{
+				message: err.Error(),
+			}
+		}
+
+		return nil, fmt.Errorf("error while creating partition consumer on partition %d, offset %d: %w", partitionID, offset, err)
+	}
+
+	defer closePartitionConsumer(pc)
+	return c.consumeSingleMessage(ctx, pc)
+}
+
+func (c durationKafkaClient) consumeSingleMessage(ctx context.Context, pc sarama.PartitionConsumer) (*sarama.ConsumerMessage, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.consumerTimeout)
+	defer cancel()
+
+	select {
+	case <-ctx.Done():
+		return nil, fmt.Errorf("error while consuming message: %w", ctx.Err())
+	case msg := <-pc.Messages():
+		return msg, nil
+	}
+}
diff --git a/component/async/kafka/simple/simple.go b/component/async/kafka/simple/simple.go
index be44ee6c6..a5cb84bfd 100644
--- a/component/async/kafka/simple/simple.go
+++ b/component/async/kafka/simple/simple.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/beatlabs/patron/component/async"
@@ -56,6 +57,7 @@ func (f *Factory) Create() (async.Consumer, error) {
 		topic:  f.topic,
 		config: cc,
 	}
+	c.partitions = c.partitionsFromOffset
 
 	for _, o := range f.oo {
 		err = o(&c.config)
@@ -64,15 +66,20 @@ func (f *Factory) Create() (async.Consumer, error) {
 		}
 	}
 
+	if c.config.DurationBasedConsumer {
+		c.partitions = c.partitionsSinceDuration
+	}
+
 	return c, nil
 }
 
 // consumer members can be injected or overwritten with the usage of OptionFunc arguments.
 type consumer struct {
-	topic  string
-	cnl    context.CancelFunc
-	ms     sarama.Consumer
-	config kafka.ConsumerConfig
+	topic      string
+	cnl        context.CancelFunc
+	ms         sarama.Consumer
+	config     kafka.ConsumerConfig
+	partitions func(context.Context) ([]sarama.PartitionConsumer, error)
 }
 
 // Close handles closing consumer.
@@ -93,10 +100,13 @@ func (c *consumer) Consume(ctx context.Context) (<-chan async.Message, <-chan er
 	chErr := make(chan error, c.config.Buffer)
 
 	log.Infof("consuming messages from topic '%s' without using consumer group", c.topic)
-	pcs, err := c.partitions()
+	var pcs []sarama.PartitionConsumer
+
+	pcs, err := c.partitions(ctx)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to get partitions: %w", err)
+		return nil, nil, fmt.Errorf("failed to get partitions since duration: %w", err)
 	}
+
 	// When kafka cluster is not fully initialized, we may get 0 partitions.
 	if len(pcs) == 0 {
 		return nil, nil, errors.New("got 0 partitions")
@@ -134,7 +144,7 @@ func (c *consumer) Consume(ctx context.Context) (<-chan async.Message, <-chan er
 	return chMsg, chErr, nil
 }
 
-func (c *consumer) partitions() ([]sarama.PartitionConsumer, error) {
+func (c *consumer) partitionsFromOffset(_ context.Context) ([]sarama.PartitionConsumer, error) {
 
 	ms, err := sarama.NewConsumer(c.config.Brokers, c.config.SaramaConfig)
 	if err != nil {
@@ -161,6 +171,56 @@ func (c *consumer) partitions() ([]sarama.PartitionConsumer, error) {
 	return pcs, nil
 }
 
+func (c *consumer) partitionsSinceDuration(ctx context.Context) ([]sarama.PartitionConsumer, error) {
+	client, err := sarama.NewClient(c.config.Brokers, c.config.SaramaConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create client: %w", err)
+	}
+
+	consumer, err := sarama.NewConsumerFromClient(client)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create simple consumer: %w", err)
+	}
+	c.ms = consumer
+
+	durationKafkaClient, err := newDurationKafkaClient(client, consumer, c.config.SaramaConfig.Net.DialTimeout)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create Kafka duration client: %w", err)
+	}
+
+	durationClient, err := newDurationClient(durationKafkaClient)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create duration client: %w", err)
+	}
+
+	offsets, err := durationClient.getTimeBasedOffsetsPerPartition(ctx, c.topic, time.Now().Add(-c.config.DurationOffset), c.config.TimeExtractor)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve duration offsets per partition: %w", err)
+	}
+
+	partitions, err := c.ms.Partitions(c.topic)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get partitions: %w", err)
+	}
+
+	pcs := make([]sarama.PartitionConsumer, len(partitions))
+
+	for i, partition := range partitions {
+		offset, exists := offsets[partition]
+		if !exists {
+			return nil, fmt.Errorf("partition %d unknown, this is most likely due to a repartitioning", partition)
+		}
+
+		pc, err := c.ms.ConsumePartition(c.topic, partition, offset)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get partition consumer: %w", err)
+		}
+		pcs[i] = pc
+	}
+
+	return pcs, nil
+}
+
 func closePartitionConsumer(cns sarama.PartitionConsumer) {
 	if cns == nil {
 		return


Test patch:
diff --git a/component/async/kafka/option_test.go b/component/async/kafka/option_test.go
index 1cc9dc5da..d9518707f 100644
--- a/component/async/kafka/option_test.go
+++ b/component/async/kafka/option_test.go
@@ -1,7 +1,9 @@
 package kafka
 
 import (
+	"errors"
 	"reflect"
+	"runtime"
 	"testing"
 	"time"
 
@@ -150,3 +152,54 @@ func TestDecoderJSON(t *testing.T) {
 		reflect.ValueOf(c.DecoderFunc).Pointer(),
 	)
 }
+
+func TestWithDurationOffset(t *testing.T) {
+	f := func(_ *sarama.ConsumerMessage) (time.Time, error) {
+		return time.Time{}, nil
+	}
+
+	type args struct {
+		since         time.Duration
+		timeExtractor TimeExtractor
+	}
+	testCases := map[string]struct {
+		args        args
+		expectedErr error
+	}{
+		"success": {
+			args: args{
+				since:         time.Second,
+				timeExtractor: f,
+			},
+		},
+		"error - negative since duration": {
+			args: args{
+				since:         -time.Second,
+				timeExtractor: f,
+			},
+			expectedErr: errors.New("duration must be positive"),
+		},
+		"error - nil time extractor": {
+			args: args{
+				since: time.Second,
+			},
+			expectedErr: errors.New("empty time extractor function"),
+		},
+	}
+	for name, tt := range testCases {
+		t.Run(name, func(t *testing.T) {
+			c := ConsumerConfig{}
+			err := WithDurationOffset(tt.args.since, tt.args.timeExtractor)(&c)
+			if tt.expectedErr != nil {
+				assert.EqualError(t, err, tt.expectedErr.Error())
+			} else {
+				assert.NoError(t, err)
+				assert.True(t, c.DurationBasedConsumer)
+				assert.Equal(t, time.Second, c.DurationOffset)
+				assert.Equal(t,
+					runtime.FuncForPC(reflect.ValueOf(tt.args.timeExtractor).Pointer()).Name(),
+					runtime.FuncForPC(reflect.ValueOf(c.TimeExtractor).Pointer()).Name())
+			}
+		})
+	}
+}
diff --git a/component/async/kafka/simple/duration_client_test.go b/component/async/kafka/simple/duration_client_test.go
new file mode 100644
index 000000000..67dedbbe0
--- /dev/null
+++ b/component/async/kafka/simple/duration_client_test.go
@@ -0,0 +1,346 @@
+package simple
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/Shopify/sarama"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Consumer_GetTimeBasedOffsetsPerPartition(t *testing.T) {
+	topic := "topic"
+	since := mustCreateTime(t, "2020-01-01T12:00:00Z")
+	dummyClient := client(topic).partitionIDs([]int32{0}, nil).build()
+	// The message is invalid as the time extractor we use require a timestamp header
+	invalidMessage := &sarama.ConsumerMessage{}
+
+	testCases := map[string]struct {
+		globalTimeout   time.Duration
+		client          *clientMock
+		expectedOffsets map[int32]int64
+		expectedErr     error
+	}{
+		"success - multiple partitions": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0, 1, 2}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since},
+						},
+					},
+				}).
+				partition(1, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						0: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-3 * time.Hour)},
+						},
+						1: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-2 * time.Hour)},
+						},
+						2: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-1 * time.Hour)},
+						},
+						3: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(1 * time.Hour)},
+						},
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(2 * time.Hour)},
+						},
+					},
+				}).
+				partition(2, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-4 * time.Hour)},
+						},
+						5: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-3 * time.Hour)},
+						},
+						6: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-2 * time.Hour)},
+						},
+						7: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-1 * time.Hour)},
+						},
+						8: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(1 * time.Hour)},
+						},
+						9: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(2 * time.Hour)},
+						},
+					},
+				}).build(),
+			expectedOffsets: map[int32]int64{
+				0: 4,
+				1: 3,
+				2: 8,
+			},
+		},
+		"success - all inside": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						0: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(1 * time.Hour)},
+						},
+						1: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(2 * time.Hour)},
+						},
+						2: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(3 * time.Hour)},
+						},
+						3: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(4 * time.Hour)},
+						},
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(5 * time.Hour)},
+						},
+					},
+				}).build(),
+			expectedOffsets: map[int32]int64{
+				0: 0,
+			},
+		},
+		"success - all outside": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-6 * time.Hour)},
+						},
+						5: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-5 * time.Hour)},
+						},
+						6: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-4 * time.Hour)},
+						},
+						7: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-3 * time.Hour)},
+						},
+						8: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-2 * time.Hour)},
+						},
+						9: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-1 * time.Hour)},
+						},
+					},
+				}).build(),
+			expectedOffsets: map[int32]int64{
+				0: 10,
+			},
+		},
+		"success - out of range offset": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since},
+							err: &outOfRangeOffsetError{
+								message: "foo",
+							},
+						},
+						5: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-3 * time.Hour)},
+						},
+						6: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-2 * time.Hour)},
+						},
+						7: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(-1 * time.Hour)},
+						},
+						8: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(1 * time.Hour)},
+						},
+						9: {
+							msg: &sarama.ConsumerMessage{Timestamp: since.Add(2 * time.Hour)},
+						},
+					},
+				}).build(),
+			expectedOffsets: map[int32]int64{
+				0: 8,
+			},
+		},
+		"error - get partitions": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs(nil, errors.New("foo")).build(),
+			expectedErr: errors.New("foo"),
+		},
+		"error - timeout": {
+			globalTimeout: time.Nanosecond,
+			client:        dummyClient,
+			expectedErr:   errors.New("context cancelled before collecting partition responses: context deadline exceeded"),
+		},
+		"error - get oldest offset": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{err: errors.New("foo")},
+				}).build(),
+			expectedErr: errors.New("foo"),
+		},
+		"error - get newest offset": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{err: errors.New("foo")},
+				}).build(),
+			expectedErr: errors.New("foo"),
+		},
+		"error - invalid message": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: invalidMessage,
+						},
+					},
+				}).build(),
+			expectedErr: errors.New("error while executing comparator: empty time"),
+		},
+		"error - get message offset": {
+			globalTimeout: time.Second,
+			client: client(topic).
+				partitionIDs([]int32{0}, nil).
+				partition(0, partitionConfig{
+					oldest: offset{offset: 0},
+					newest: offset{offset: 10},
+					messages: map[int64]messageAtOffset{
+						4: {
+							msg: &sarama.ConsumerMessage{Timestamp: since},
+							err: errors.New("foo"),
+						},
+					},
+				}).build(),
+			expectedErr: errors.New("error while retrieving message offset 4 on partition 0: foo"),
+		},
+	}
+	for name, tt := range testCases {
+		t.Run(name, func(t *testing.T) {
+			consumer, err := newDurationClient(tt.client)
+			require.NoError(t, err)
+			ctx, cancel := context.WithTimeout(context.Background(), tt.globalTimeout)
+			defer cancel()
+
+			offsets, err := consumer.getTimeBasedOffsetsPerPartition(ctx, topic, since, kafkaHeaderTimeExtractor)
+			if tt.expectedErr != nil {
+				assert.EqualError(t, err, tt.expectedErr.Error())
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expectedOffsets, offsets)
+			}
+		})
+	}
+}
+
+type clientBuilder struct {
+	topic            string
+	partitions       []int32
+	partitionsError  error
+	partitionConfigs map[int32]partitionConfig
+}
+
+func client(topic string) *clientBuilder {
+	return &clientBuilder{
+		topic:            topic,
+		partitionConfigs: make(map[int32]partitionConfig),
+	}
+}
+
+func (c *clientBuilder) partitionIDs(partitions []int32, err error) *clientBuilder {
+	c.partitions = partitions
+	c.partitionsError = err
+	return c
+}
+
+type partitionConfig struct {
+	oldest   offset
+	newest   offset
+	messages map[int64]messageAtOffset
+}
+
+type offset struct {
+	offset int64
+	err    error
+}
+
+type messageAtOffset struct {
+	msg *sarama.ConsumerMessage
+	err error
+}
+
+func (c *clientBuilder) partition(partitionID int32, partitionConfig partitionConfig) *clientBuilder {
+	c.partitionConfigs[partitionID] = partitionConfig
+	return c
+}
+
+func (c *clientBuilder) build() *clientMock {
+	return &clientMock{
+		builder: c,
+	}
+}
+
+type clientMock struct {
+	builder *clientBuilder
+}
+
+func (c *clientMock) getPartitionIDs(_ string) ([]int32, error) {
+	return c.builder.partitions, c.builder.partitionsError
+}
+
+func (c *clientMock) getOldestOffset(_ string, partitionID int32) (int64, error) {
+	return c.builder.partitionConfigs[partitionID].oldest.offset, c.builder.partitionConfigs[partitionID].oldest.err
+}
+
+func (c *clientMock) getNewestOffset(_ string, partitionID int32) (int64, error) {
+	return c.builder.partitionConfigs[partitionID].newest.offset, c.builder.partitionConfigs[partitionID].newest.err
+}
+
+func (c *clientMock) getMessageAtOffset(_ context.Context, _ string, partitionID int32, offset int64) (*sarama.ConsumerMessage, error) {
+	cfg := c.builder.partitionConfigs[partitionID].messages[offset]
+	return cfg.msg, cfg.err
+}
+
+func mustCreateTime(t *testing.T, timestamp string) time.Time {
+	ts, err := time.Parse(time.RFC3339, timestamp)
+	require.NoError(t, err)
+	return ts
+}
+
+func kafkaHeaderTimeExtractor(msg *sarama.ConsumerMessage) (time.Time, error) {
+	if msg.Timestamp.IsZero() {
+		return time.Time{}, errors.New("empty time")
+	}
+	return msg.Timestamp, nil
+}
diff --git a/test/docker/kafka/consumer_simple_integration_test.go b/test/docker/kafka/consumer_simple_integration_test.go
index 6f476c4ab..44929515d 100644
--- a/test/docker/kafka/consumer_simple_integration_test.go
+++ b/test/docker/kafka/consumer_simple_integration_test.go
@@ -107,3 +107,74 @@ func TestSimpleConsume_ClaimMessageError(t *testing.T) {
 		require.EqualError(t, err, "could not determine decoder  failed to determine content type from message headers [] : content type header is missing")
 	}
 }
+
+func TestSimpleConsume_WithDurationOffset(t *testing.T) {
+	now := time.Now()
+	sent := createTimestampPayload(
+		now.Add(-10*time.Hour),
+		now.Add(-5*time.Hour),
+		now.Add(-3*time.Hour),
+		now.Add(-2*time.Hour),
+		now.Add(-1*time.Hour),
+	)
+
+	messages := make([]*sarama.ProducerMessage, 0)
+	for _, val := range sent {
+		messages = append(messages, getProducerMessage(simpleTopic3, val))
+	}
+
+	err := sendMessages(messages...)
+	require.NoError(t, err)
+
+	chMessages := make(chan []string)
+	chErr := make(chan error)
+	go func() {
+		factory, err := simple.New("test1", simpleTopic3, Brokers(), kafka.DecoderJSON(), kafka.Version(sarama.V2_1_0_0.String()),
+			kafka.StartFromNewest(), kafka.WithDurationOffset(4*time.Hour, timestampExtractor))
+		if err != nil {
+			chErr <- err
+			return
+		}
+
+		consumer, err := factory.Create()
+		if err != nil {
+			chErr <- err
+			return
+		}
+		defer func() {
+			_ = consumer.Close()
+		}()
+
+		received, err := consumeMessages(consumer, 3)
+		if err != nil {
+			chErr <- err
+			return
+		}
+
+		chMessages <- received
+	}()
+
+	time.Sleep(5 * time.Second)
+
+	var received []string
+
+	select {
+	case received = <-chMessages:
+	case err = <-chErr:
+		require.NoError(t, err)
+	}
+
+	assert.Equal(t, sent[2:], received)
+}
+
+func createTimestampPayload(timestamps ...time.Time) []string {
+	payloads := make([]string, len(timestamps))
+	for i, timestamp := range timestamps {
+		payloads[i] = timestamp.Format(time.RFC3339)
+	}
+	return payloads
+}
+
+func timestampExtractor(msg *sarama.ConsumerMessage) (time.Time, error) {
+	return time.Parse(time.RFC3339, string(msg.Value))
+}
diff --git a/test/docker/kafka/kafka_test.go b/test/docker/kafka/kafka_test.go
index 5d23b8642..29775f6a3 100644
--- a/test/docker/kafka/kafka_test.go
+++ b/test/docker/kafka/kafka_test.go
@@ -25,6 +25,7 @@ const (
 	zookeeperPort = "2181"
 	simpleTopic1  = "simpleTopic1"
 	simpleTopic2  = "simpleTopic2"
+	simpleTopic3  = "simpleTopic3"
 	groupTopic1   = "groupTopic1"
 	groupTopic2   = "groupTopic2"
 )
@@ -188,6 +189,8 @@ func consumeMessages(consumer async.Consumer, expectedMessageCount int) ([]strin
 
 	for {
 		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
 		case msg := <-ch:
 			received = append(received, string(msg.Payload()))
 			expectedMessageCount--