diff --git a/consumer_test.go b/consumer_test.go index 2ec2ec4..c68cb17 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -3,6 +3,7 @@ package consumer import ( "context" "fmt" + "github.com/harlow/kinesis-consumer/store/memory" "sync" "testing" @@ -52,7 +53,7 @@ func TestScan(t *testing.T) { }, } var ( - cp = &fakeCheckpoint{cache: map[string]string{}} + cp = memory.New() ctr = &fakeCounter{} ) @@ -114,7 +115,7 @@ func TestScanShard(t *testing.T) { } var ( - cp = &fakeCheckpoint{cache: map[string]string{}} + cp = memory.New() ctr = &fakeCounter{} ) @@ -219,7 +220,7 @@ func TestScanShard_SkipCheckpoint(t *testing.T) { }, } - var cp = &fakeCheckpoint{cache: map[string]string{}} + var cp = memory.New() c, err := New("myStreamName", WithClient(client), WithStore(cp)) if err != nil { @@ -335,29 +336,6 @@ func (c *kinesisClientMock) GetShardIteratorWithContext(ctx aws.Context, in *kin return c.getShardIteratorMock(in) } -// implementation of checkpoint -type fakeCheckpoint struct { - cache map[string]string - mu sync.Mutex -} - -func (fc *fakeCheckpoint) SetCheckpoint(streamName, shardID, sequenceNumber string) error { - fc.mu.Lock() - defer fc.mu.Unlock() - - key := fmt.Sprintf("%s-%s", streamName, shardID) - fc.cache[key] = sequenceNumber - return nil -} - -func (fc *fakeCheckpoint) GetCheckpoint(streamName, shardID string) (string, error) { - fc.mu.Lock() - defer fc.mu.Unlock() - - key := fmt.Sprintf("%s-%s", streamName, shardID) - return fc.cache[key], nil -} - // implementation of counter type fakeCounter struct { counter int64 diff --git a/store/memory/store.go b/store/memory/store.go new file mode 100644 index 0000000..a556113 --- /dev/null +++ b/store/memory/store.go @@ -0,0 +1,33 @@ +// The memory store provides a store that can be used for testing and single-threaded applications. +// DO NOT USE this in a production application where persistence beyond a single application lifecycle is necessary +// or when there are multiple consumers. +package memory + +import ( + "fmt" + "sync" +) + +func New() *Checkpoint{ + return &Checkpoint{} +} + +type Checkpoint struct { + sync.Map +} + +func (c *Checkpoint) SetCheckpoint(streamName, shardID, sequenceNumber string) error { + if sequenceNumber == "" { + return fmt.Errorf("sequence number should not be empty") + } + c.Store(streamName+":"+shardID, sequenceNumber) + return nil +} + +func (c *Checkpoint) GetCheckpoint(streamName, shardID string) (string, error) { + val, ok := c.Load(streamName + ":" + shardID) + if !ok { + return "", nil + } + return val.(string), nil +} diff --git a/store/memory/store_test.go b/store/memory/store_test.go new file mode 100644 index 0000000..6e2041c --- /dev/null +++ b/store/memory/store_test.go @@ -0,0 +1,32 @@ +package memory_test + +import ( + "testing" + + "github.com/harlow/kinesis-consumer/store/memory" +) + +func Test_CheckpointLifecycle(t *testing.T) { + c := memory.New() + + // set + c.SetCheckpoint("streamName", "shardID", "testSeqNum") + + // get + val, err := c.GetCheckpoint("streamName", "shardID") + if err != nil { + t.Fatalf("get checkpoint error: %v", err) + } + if val != "testSeqNum" { + t.Fatalf("checkpoint exists expected %s, got %s", "testSeqNum", val) + } +} + +func Test_SetEmptySeqNum(t *testing.T) { + c := memory.New() + + err := c.SetCheckpoint("streamName", "shardID", "") + if err == nil || err.Error() != "sequence number should not be empty" { + t.Fatalf("should not allow empty sequence number") + } +}