Skip to content
This repository has been archived by the owner on Jan 8, 2020. It is now read-only.

Expose offset methods on partition consumers #209

Merged
merged 2 commits into from Mar 7, 2018
Merged

Expose offset methods on partition consumers #209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
895dd20
Expose offset methods on partition consumers
dim Jan 16, 2018
60c9f44
Fix/remove test
dim Mar 7, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 0 additions & 4 deletions balancer.go
View file
Open in desktop
Expand Up @@ -157,10 +157,6 @@ func (r *balancer) Topic(name string, memberID string) error {
}

func (r *balancer) Perform(s Strategy) map[string]map[string][]int32 {
if r == nil {
return nil
}

res := make(map[string]map[string][]int32, 1)
for topic, info := range r.topics {
for memberID, partitions := range info.Perform(s) {
Expand Down
48 changes: 22 additions & 26 deletions consumer.go
View file
Open in desktop
Expand Up @@ -140,18 +140,16 @@ func (c *Consumer) HighWaterMarks() map[string]map[int32]int64 { return c.consum
// your application crashes. This means that you may end up processing the same
// message twice, and your processing should ideally be idempotent.
func (c *Consumer) MarkOffset(msg *sarama.ConsumerMessage, metadata string) {
sub := c.subs.Fetch(msg.Topic, msg.Partition)
if sub != nil {
sub.MarkOffset(msg.Offset+1, metadata)
if sub := c.subs.Fetch(msg.Topic, msg.Partition); sub != nil {
sub.MarkOffset(msg.Offset, metadata)
}
}

// MarkPartitionOffset marks an offset of the provided topic/partition as processed.
// See MarkOffset for additional explanation.
func (c *Consumer) MarkPartitionOffset(topic string, partition int32, offset int64, metadata string) {
sub := c.subs.Fetch(topic, partition)
if sub != nil {
sub.MarkOffset(offset+1, metadata)
if sub := c.subs.Fetch(topic, partition); sub != nil {
sub.MarkOffset(offset, metadata)
}
}

Expand All @@ -162,9 +160,8 @@ func (c *Consumer) MarkOffsets(s *OffsetStash) {
defer s.mu.Unlock()

for tp, info := range s.offsets {
sub := c.subs.Fetch(tp.Topic, tp.Partition)
if sub != nil {
sub.MarkOffset(info.Offset+1, info.Metadata)
if sub := c.subs.Fetch(tp.Topic, tp.Partition); sub != nil {
sub.MarkOffset(info.Offset, info.Metadata)
}
delete(s.offsets, tp)
}
Expand All @@ -177,9 +174,8 @@ func (c *Consumer) MarkOffsets(s *OffsetStash) {
//
// Difference between ResetOffset and MarkOffset is that it allows to rewind to an earlier offset
func (c *Consumer) ResetOffset(msg *sarama.ConsumerMessage, metadata string) {
sub := c.subs.Fetch(msg.Topic, msg.Partition)
if sub != nil {
sub.ResetOffset(msg.Offset+1, metadata)
if sub := c.subs.Fetch(msg.Topic, msg.Partition); sub != nil {
sub.ResetOffset(msg.Offset, metadata)
}
}

Expand All @@ -188,7 +184,7 @@ func (c *Consumer) ResetOffset(msg *sarama.ConsumerMessage, metadata string) {
func (c *Consumer) ResetPartitionOffset(topic string, partition int32, offset int64, metadata string) {
sub := c.subs.Fetch(topic, partition)
if sub != nil {
sub.ResetOffset(offset+1, metadata)
sub.ResetOffset(offset, metadata)
}
}

Expand All @@ -199,9 +195,8 @@ func (c *Consumer) ResetOffsets(s *OffsetStash) {
defer s.mu.Unlock()

for tp, info := range s.offsets {
sub := c.subs.Fetch(tp.Topic, tp.Partition)
if sub != nil {
sub.ResetOffset(info.Offset+1, info.Metadata)
if sub := c.subs.Fetch(tp.Topic, tp.Partition); sub != nil {
sub.ResetOffset(info.Offset, info.Metadata)
}
delete(s.offsets, tp)
}
Expand Down Expand Up @@ -264,9 +259,8 @@ func (c *Consumer) CommitOffsets() error {
if kerr != sarama.ErrNoError {
err = kerr
} else if state, ok := snap[topicPartition{topic, partition}]; ok {
sub := c.subs.Fetch(topic, partition)
if sub != nil {
sub.MarkCommitted(state.Info.Offset)
if sub := c.subs.Fetch(topic, partition); sub != nil {
sub.markCommitted(state.Info.Offset)
}
}
}
Expand Down Expand Up @@ -706,11 +700,13 @@ func (c *Consumer) syncGroup(strategy *balancer) (map[string][]int32, error) {
GenerationId: generationID,
}

for memberID, topics := range strategy.Perform(c.client.config.Group.PartitionStrategy) {
if err := req.AddGroupAssignmentMember(memberID, &sarama.ConsumerGroupMemberAssignment{
Topics: topics,
}); err != nil {
return nil, err
if strategy != nil {
for memberID, topics := range strategy.Perform(c.client.config.Group.PartitionStrategy) {
if err := req.AddGroupAssignmentMember(memberID, &sarama.ConsumerGroupMemberAssignment{
Topics: topics,
}); err != nil {
return nil, err
}
}
}

Expand Down Expand Up @@ -828,9 +824,9 @@ func (c *Consumer) createConsumer(tomb *loopTomb, topic string, partition int32,
// Start partition consumer goroutine
tomb.Go(func(stopper <-chan none) {
if c.client.config.Group.Mode == ConsumerModePartitions {
pc.WaitFor(stopper, c.errors)
pc.waitFor(stopper, c.errors)
} else {
pc.Multiplex(stopper, c.messages, c.errors)
pc.multiplex(stopper, c.messages, c.errors)
}
})

Expand Down
42 changes: 17 additions & 25 deletions partitions.go
View file
Open in desktop
Expand Up @@ -19,6 +19,12 @@ type PartitionConsumer interface {

// Partition returns the consumed partition
Partition() int32

// MarkOffset marks the offset of a message as preocessed.
MarkOffset(offset int64, metadata string)

// ResetOffset resets the offset to a previously processed message.
ResetOffset(offset int64, metadata string)
}

type partitionConsumer struct {
Expand Down Expand Up @@ -81,7 +87,7 @@ func (c *partitionConsumer) Close() error {
return c.closeErr
}

func (c *partitionConsumer) WaitFor(stopper <-chan none, errors chan<- error) {
func (c *partitionConsumer) waitFor(stopper <-chan none, errors chan<- error) {
defer close(c.dead)

for {
Expand All @@ -105,7 +111,7 @@ func (c *partitionConsumer) WaitFor(stopper <-chan none, errors chan<- error) {
}
}

func (c *partitionConsumer) Multiplex(stopper <-chan none, messages chan<- *sarama.ConsumerMessage, errors chan<- error) {
func (c *partitionConsumer) multiplex(stopper <-chan none, messages chan<- *sarama.ConsumerMessage, errors chan<- error) {
defer close(c.dead)

for {
Expand Down Expand Up @@ -140,52 +146,38 @@ func (c *partitionConsumer) Multiplex(stopper <-chan none, messages chan<- *sara
}
}

func (c *partitionConsumer) State() partitionState {
if c == nil {
return partitionState{}
}

func (c *partitionConsumer) getState() partitionState {
c.mu.Lock()
state := c.state
c.mu.Unlock()

return state
}

func (c *partitionConsumer) MarkCommitted(offset int64) {
if c == nil {
return
}

func (c *partitionConsumer) markCommitted(offset int64) {
c.mu.Lock()
if offset == c.state.Info.Offset {
c.state.Dirty = false
}
c.mu.Unlock()
}

// MarkOffset implements PartitionConsumer
func (c *partitionConsumer) MarkOffset(offset int64, metadata string) {
if c == nil {
return
}

c.mu.Lock()
if offset > c.state.Info.Offset {
c.state.Info.Offset = offset
if next := offset + 1; next > c.state.Info.Offset {
c.state.Info.Offset = next
c.state.Info.Metadata = metadata
c.state.Dirty = true
}
c.mu.Unlock()
}

// ResetOffset implements PartitionConsumer
func (c *partitionConsumer) ResetOffset(offset int64, metadata string) {
if c == nil {
return
}

c.mu.Lock()
if offset <= c.state.Info.Offset {
c.state.Info.Offset = offset
if next := offset + 1; next <= c.state.Info.Offset {
c.state.Info.Offset = next
c.state.Info.Metadata = metadata
c.state.Dirty = true
}
Expand Down Expand Up @@ -244,7 +236,7 @@ func (m *partitionMap) Snapshot() map[topicPartition]partitionState {

snap := make(map[topicPartition]partitionState, len(m.data))
for tp, pc := range m.data {
snap[tp] = pc.State()
snap[tp] = pc.getState()
}
return snap
}
Expand Down
53 changes: 22 additions & 31 deletions partitions_test.go
View file
Open in desktop
Expand Up @@ -21,7 +21,7 @@ var _ = Describe("partitionConsumer", func() {
})

It("should set state", func() {
Expect(subject.State()).To(Equal(partitionState{
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2000, "m3ta"},
}))
})
Expand All @@ -32,67 +32,58 @@ var _ = Describe("partitionConsumer", func() {
defer pc.Close()
close(pc.dead)

state := pc.State()
state := pc.getState()
Expect(state.Info.Offset).To(Equal(int64(-1)))
Expect(state.Info.Metadata).To(Equal("m3ta"))
})

It("should update state", func() {
subject.MarkOffset(2001, "met@") // should set state
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2001, "met@"},
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2002, "met@"},
Dirty: true,
}))

subject.MarkCommitted(2001) // should reset dirty status
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2001, "met@"},
subject.markCommitted(2002) // should reset dirty status
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2002, "met@"},
}))

subject.MarkOffset(2001, "me7a") // should not update state
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2001, "met@"},
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2002, "met@"},
}))

subject.MarkOffset(2002, "me7a") // should bump state
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2002, "me7a"},
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2003, "me7a"},
Dirty: true,
}))

// After committing a later offset, try rewinding back to earlier offset with new metadata.
subject.ResetOffset(2001, "met@")
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2001, "met@"},
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2002, "met@"},
Dirty: true,
}))

subject.MarkCommitted(2001) // should not unset state
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2001, "met@"},
subject.markCommitted(2002) // should not unset state
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2002, "met@"},
}))

subject.MarkOffset(2002, "me7a") // should bump state
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2002, "me7a"},
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2003, "me7a"},
Dirty: true,
}))

subject.MarkCommitted(2002)
Expect(subject.State()).To(Equal(partitionState{
Info: offsetInfo{2002, "me7a"},
subject.markCommitted(2003)
Expect(subject.getState()).To(Equal(partitionState{
Info: offsetInfo{2003, "me7a"},
}))
})

It("should not fail when nil", func() {
blank := (*partitionConsumer)(nil)
Expect(func() {
_ = blank.State()
blank.MarkOffset(2001, "met@")
blank.MarkCommitted(2001)
}).NotTo(Panic())
})

})

var _ = Describe("partitionMap", func() {
Expand Down Expand Up @@ -135,7 +126,7 @@ var _ = Describe("partitionMap", func() {

subject.Store("topic", 0, pc0)
subject.Store("topic", 1, pc1)
subject.Fetch("topic", 1).MarkOffset(2001, "met@")
subject.Fetch("topic", 1).MarkOffset(2000, "met@")

Expect(subject.Snapshot()).To(Equal(map[topicPartition]partitionState{
{"topic", 0}: {Info: offsetInfo{2000, "m3ta"}, Dirty: false},
Expand Down