wip

offset_fetch_response.go

@@ -80,3 +80,15 @@ func (r *OffsetFetchResponse) decode(pd packetDecoder) (err error) {
 
 	return nil
 }
+
+func (r *OffsetFetchResponse) GetBlock(topic string, partition int32) *OffsetFetchResponseBlock {
+	if r.Blocks == nil {
+		return nil
+	}
+
+	if r.Blocks[topic] == nil {
+		return nil
+	}
+
+	return r.Blocks[topic][partition]
+}

offset_manager.go

@@ -0,0 +1,369 @@
+package sarama
+
+import (
+	"sync"
+	"time"
+)
+
+// Offset Manager
+
+// OffsetManager uses Kafka to store and fetch consumed partition offsets.
+type OffsetManager interface {
+	// ManagePartition creates a PartitionOffsetManager on the given topic/partition. It will
+	// return an error if this OffsetManager is already managing the given topic/partition.
+	ManagePartition(topic string, partition int32) (PartitionOffsetManager, error)
+}
+
+type offsetManager struct {
+	client Client
+	conf   *Config
+	group  string
+
+	lock sync.Mutex
+	poms map[string]map[int32]*partitionOffsetManager
+	boms map[*Broker]*brokerOffsetManager
+}
+
+// NewOffsetManagerFromClient creates a new OffsetManager from the given client.
+// It is still necessary to call Close() on the underlying client when finished with the partition manager.
+func NewOffsetManagerFromClient(client Client) (OffsetManager, error) {
+	// Check that we are not dealing with a closed Client before processing any other arguments
+	if client.Closed() {
+		return nil, ErrClosedClient
+	}
+
+	om := &offsetManager{
+		client: client,
+		conf:   client.Config(),
+		poms:   make(map[string]map[int32]*partitionOffsetManager),
+		boms:   make(map[*Broker]*brokerOffsetManager),
+	}
+
+	return om, nil
+}
+
+func (om *offsetManager) ManagePartition(topic string, partition int32) (PartitionOffsetManager, error) {
+	pom, err := om.newPartitionOffsetManager(topic, partition)
+	if err != nil {
+		return nil, err
+	}
+
+	om.lock.Lock()
+	defer om.lock.Unlock()
+
+	topicManagers := om.poms[topic]
+	if topicManagers == nil {
+		topicManagers = make(map[int32]*partitionOffsetManager)
+		om.poms[topic] = topicManagers
+	}
+
+	if topicManagers[partition] != nil {
+		return nil, ConfigurationError("That topic/partition is already being managed")
+	}
+
+	topicManagers[partition] = pom
+	return pom, nil
+}
+
+func (om *offsetManager) refBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
+	om.lock.Lock()
+	defer om.lock.Unlock()
+
+	bom := om.boms[broker]
+	if bom == nil {
+		bom = om.newBrokerOffsetManager(broker)
+		om.boms[broker] = bom
+	}
+
+	bom.refs++
+
+	return bom
+}
+
+func (om *offsetManager) unrefBrokerOffsetManager(bom *brokerOffsetManager) {
+	om.lock.Lock()
+	defer om.lock.Unlock()
+
+	bom.refs--
+
+	if bom.refs == 0 {
+		close(bom.newSubscriptions)
+		if om.boms[bom.broker] == bom {
+			delete(om.boms, bom.broker)
+		}
+	}
+}
+
+func (om *offsetManager) abandonBroker(bom *brokerOffsetManager) {
+	om.lock.Lock()
+	defer om.lock.Unlock()
+
+	delete(om.boms, bom.broker)
+}
+
+// Partition Offset Manager
+
+// PartitionOffsetManager uses Kafka to store and fetch consumed partition offsets. You MUST call Close()
+// on a partition offset manager to avoid leaks, it will not be garbage-collected automatically when it passes
+// out of scope.
+type PartitionOffsetManager interface {
+	// TODO these need documentation
+	Offset() int64
+	SetOffset(offset int64)
+
+	Metadata() string
+	SetMetadata(metadata string)
+
+	Errors() <-chan error
+	AsyncClose()
+	Close() error
+}
+
+type partitionOffsetManager struct {
+	parent    *offsetManager
+	topic     string
+	partition int32
+
+	lock     sync.Mutex
+	offset   int64
+	metadata string
+	broker   *brokerOffsetManager
+
+	errors    chan error
+	rebalance chan none
+}
+
+func (om *offsetManager) newPartitionOffsetManager(topic string, partition int32) (*partitionOffsetManager, error) {
+	pom := &partitionOffsetManager{
+		parent:    om,
+		topic:     topic,
+		partition: partition,
+		errors:    make(chan error, om.conf.ChannelBufferSize),
+		rebalance: make(chan none, 1),
+	}
+
+	if err := pom.selectBroker(); err != nil {
+		return nil, err
+	}
+
+	if err := pom.fetchInitialOffset(); err != nil {
+		return nil, err
+	}
+
+	go withRecover(pom.mainLoop)
+
+	return pom, nil
+}
+
+func (pom *partitionOffsetManager) mainLoop() {
+	for _ = range pom.rebalance {
+		if err := pom.selectBroker(); err != nil {
+			pom.handleError(err)
+			pom.rebalance <- none{}
+		}
+	}
+}
+
+func (pom *partitionOffsetManager) selectBroker() error {
+	if pom.broker != nil {
+		pom.parent.unrefBrokerOffsetManager(pom.broker)
+		pom.broker = nil
+	}
+
+	var broker *Broker
+	var err error
+
+	if err = pom.parent.client.RefreshCoordinator(pom.parent.group); err != nil {
+		return err
+	}
+
+	if broker, err = pom.parent.client.Coordinator(pom.parent.group); err != nil {
+		return err
+	}
+
+	pom.broker = pom.parent.refBrokerOffsetManager(broker)
+	pom.broker.newSubscriptions <- pom
+	return nil
+}
+
+func (pom *partitionOffsetManager) fetchInitialOffset() error {
+	request := new(OffsetFetchRequest)
+	request.Version = 1 // TODO should the version be configurable?
+	request.ConsumerGroup = pom.parent.group
+	request.AddPartition(pom.topic, pom.partition)
+
+	response, err := pom.broker.broker.FetchOffset(request)
+	if err != nil {
+		return err
+	}
+
+	block := response.GetBlock(pom.topic, pom.partition)
+	if block == nil {
+		return ErrIncompleteResponse
+	}
+
+	switch block.Err {
+	case ErrNoError:
+		pom.offset = block.Offset
+		pom.metadata = block.Metadata
+		return nil
+	default:
+		// TODO what other errors can occur here, should we retry some of them?
+		return err
+	}
+}
+
+func (pom *partitionOffsetManager) handleError(err error) {
+	// TODO should OffsetManager have its own section of `Config` or should it just borrow Consumer's?
+	if pom.parent.conf.Consumer.Return.Errors {
+		pom.errors <- err
+	} else {
+		Logger.Println(err)
+	}
+}
+
+func (pom *partitionOffsetManager) Errors() <-chan error {
+	return pom.errors
+}
+
+func (pom *partitionOffsetManager) SetOffset(offset int64) {
+	pom.lock.Lock()
+	defer pom.lock.Unlock()
+
+	pom.offset = offset
+}
+
+func (pom *partitionOffsetManager) SetMetadata(metadata string) {
+	pom.lock.Lock()
+	defer pom.lock.Unlock()
+
+	pom.metadata = metadata
+}
+
+func (pom *partitionOffsetManager) Offset() int64 {
+	pom.lock.Lock()
+	defer pom.lock.Unlock()
+
+	return pom.offset
+}
+
+func (pom *partitionOffsetManager) Metadata() string {
+	pom.lock.Lock()
+	defer pom.lock.Unlock()
+
+	return pom.metadata
+}
+
+func (pom *partitionOffsetManager) AsyncClose() {
+	// TODO implement me
+}
+
+func (pom *partitionOffsetManager) Close() error {
+	pom.AsyncClose()
+	// TODO read from errors (do we need another error type a la ConsumerError(s)?)
+	return nil
+}
+
+// Broker Offset Manager
+
+type brokerOffsetManager struct {
+	parent           *offsetManager
+	broker           *Broker
+	timer            *time.Ticker
+	newSubscriptions chan *partitionOffsetManager
+	subscriptions    map[*partitionOffsetManager]none
+	refs             int
+}
+
+func (om *offsetManager) newBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
+	bom := &brokerOffsetManager{
+		parent:           om,
+		broker:           broker,
+		timer:            time.NewTicker(5 * time.Second), // TODO this should be configurable
+		newSubscriptions: make(chan *partitionOffsetManager),
+		subscriptions:    make(map[*partitionOffsetManager]none),
+	}
+
+	go withRecover(bom.mainLoop)
+
+	return bom
+}
+
+func (bom *brokerOffsetManager) mainLoop() {
+	for {
+		select {
+		case <-bom.timer.C:
+			bom.flushToBroker()
+		case s, ok := <-bom.newSubscriptions:
+			if !ok {
+				bom.timer.Stop()
+				return
+			}
+			bom.subscriptions[s] = none{}
+		}
+	}
+}
+
+func (bom *brokerOffsetManager) flushToBroker() {
+	request := bom.constructRequest()
+	response, err := bom.broker.CommitOffset(request)
+
+	if err != nil {
+		bom.abort(err)
+	}
+
+	for s := range bom.subscriptions {
+		var err KError
+		var ok bool
+
+		if response.Errors[s.topic] == nil {
+			s.handleError(ErrIncompleteResponse)
+			delete(bom.subscriptions, s)
+			s.rebalance <- none{}
+			continue
+		}
+		if err, ok = response.Errors[s.topic][s.partition]; !ok {
+			s.handleError(ErrIncompleteResponse)
+			delete(bom.subscriptions, s)
+			s.rebalance <- none{}
+			continue
+		}
+
+		switch err {
+		case ErrNoError:
+			break
+		case ErrUnknownTopicOrPartition, ErrNotLeaderForPartition, ErrLeaderNotAvailable:
+			delete(bom.subscriptions, s)
+			s.rebalance <- none{}
+		default:
+			s.handleError(err)
+			delete(bom.subscriptions, s)
+			s.rebalance <- none{}
+		}
+	}
+}
+
+func (bom *brokerOffsetManager) constructRequest() *OffsetCommitRequest {
+	r := &OffsetCommitRequest{} // TODO use the right version
+	for s := range bom.subscriptions {
+		s.lock.Lock()
+		r.AddBlock(s.topic, s.partition, s.offset, 0, s.metadata)
+		s.lock.Unlock()
+	}
+	return r
+}
+
+func (bom *brokerOffsetManager) abort(err error) {
+	bom.parent.abandonBroker(bom)
+	_ = bom.broker.Close() // we don't care about the error this might return, we already have one
+
+	for pom := range bom.subscriptions {
+		pom.handleError(err)
+		pom.rebalance <- none{}
+	}
+
+	for s := range bom.newSubscriptions {
+		s.handleError(err)
+		s.rebalance <- none{}
+	}
+}