[KafkaSource] Use the control protocol to expose the consumer group status

go.mod

@@ -33,6 +33,7 @@ require (
  k8s.io/apimachinery v0.19.7
  k8s.io/client-go v0.19.7
  k8s.io/utils v0.0.0-20200729134348-d5654de09c73
+ knative.dev/control-protocol v0.0.0-20210402065223-6e2383c7197d
  knative.dev/eventing v0.22.1-0.20210406191848-5171353bd1ed
  knative.dev/hack v0.0.0-20210325223819-b6ab329907d3
  knative.dev/networking v0.0.0-20210406043338-e38eb2be3962

go.sum

@@ -1276,6 +1276,8 @@ k8s.io/kube-openapi v0.0.0-20200805222855-6aeccd4b50c6 h1:+WnxoVtG8TMiudHBSEtrVL
 k8s.io/kube-openapi v0.0.0-20200805222855-6aeccd4b50c6/go.mod h1:UuqjUnNftUyPE5H64/qeyjQoUZhGpeFDVdxjTeEVN2o=
 k8s.io/utils v0.0.0-20200729134348-d5654de09c73 h1:uJmqzgNWG7XyClnU/mLPBWwfKKF1K8Hf8whTseBgJcg=
 k8s.io/utils v0.0.0-20200729134348-d5654de09c73/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
+knative.dev/control-protocol v0.0.0-20210402065223-6e2383c7197d h1:1ZueD3ovMOj8d/Ez+uPTol9jCwgy/6vkcGj26i3I498=
+knative.dev/control-protocol v0.0.0-20210402065223-6e2383c7197d/go.mod h1:82/EwanWb1EjpPVz7ts5IxLi0GlX0eFZ5oVa8/xLMuc=
 knative.dev/eventing v0.22.1-0.20210406191848-5171353bd1ed h1:TDGiR+w4hHN+VaXWF4bg8mb7Avwen+pjPfwSJbVQbxA=
 knative.dev/eventing v0.22.1-0.20210406191848-5171353bd1ed/go.mod h1:LOG7bh0eZQkbYANcnORwke6Yy6aUu62o8GeByaOFfRQ=
 knative.dev/hack v0.0.0-20210325223819-b6ab329907d3 h1:km0Rrh0T9/wA2pivQm1hqSPVwgNgGCHC2WNn3GakZmE=
@@ -1286,6 +1288,7 @@ knative.dev/networking v0.0.0-20210406043338-e38eb2be3962/go.mod h1:0V6M1AaWPL/M
 knative.dev/pkg v0.0.0-20210329065222-9d92ea16c0d3/go.mod h1:PD5g8hUCXq6iR3tILjmZeJBvQfXGnHMPKryq54qHJhg=
 knative.dev/pkg v0.0.0-20210330162221-808d62257db6/go.mod h1:PD5g8hUCXq6iR3tILjmZeJBvQfXGnHMPKryq54qHJhg=
 knative.dev/pkg v0.0.0-20210331062321-6317ec6066f4/go.mod h1:PD5g8hUCXq6iR3tILjmZeJBvQfXGnHMPKryq54qHJhg=
+knative.dev/pkg v0.0.0-20210331065221-952fdd90dbb0/go.mod h1:PD5g8hUCXq6iR3tILjmZeJBvQfXGnHMPKryq54qHJhg=
 knative.dev/pkg v0.0.0-20210406170139-b8e331a6abf3 h1:mm1fYtGW1NPYbX5+YtvXx8v4Hjqil7Vb1xc6nHpTDEY=
 knative.dev/pkg v0.0.0-20210406170139-b8e331a6abf3/go.mod h1:PD5g8hUCXq6iR3tILjmZeJBvQfXGnHMPKryq54qHJhg=
 knative.dev/reconciler-test v0.0.0-20210329214021-2a67496462a4/go.mod h1:qLfJMHc0i9ENTSet/SUp/FcQm4QVfNTX8ZC//aVQN0M=

pkg/apis/sources/v1alpha1/kafka_conversion.go

@@ -47,6 +47,7 @@ func (source *KafkaSource) ConvertTo(ctx context.Context, obj apis.Convertible)
  }
  source.Status.Status.DeepCopyInto(&sink.Status.Status)
  sink.Status.Consumers = source.Status.Consumers
+ sink.Status.Claims = source.Status.Claims
  sink.Status.Selector = source.Status.Selector
  source.Status.Placeable.DeepCopyInto(&sink.Status.Placeable)
 
@@ -97,6 +98,7 @@ func (sink *KafkaSource) ConvertFrom(ctx context.Context, obj apis.Convertible)
  source.Status.Status.DeepCopyInto(&sink.Status.Status)
  sink.Status.Consumers = source.Status.Consumers
  sink.Status.Selector = source.Status.Selector
+ sink.Status.Claims = source.Status.Claims
  source.Status.Placeable.DeepCopyInto(&sink.Status.Placeable)
 
  // Optionals

pkg/apis/sources/v1alpha1/kafka_types.go

@@ -137,6 +137,11 @@ type KafkaSourceStatus struct {
  // For round-tripping only.
  Consumers int32 `json:"consumers,omitempty"`
 
+ // Claims consumed by this KafkaSource instance
+ // +optional
+ // For round-tripping only.
+ Claims string `json:"claims,omitempty"`
+
  // Use for labelSelectorPath when scaling Kafka source
  // +optional
  Selector string `json:"selector,omitempty"`

pkg/apis/sources/v1beta1/kafka_lifecycle.go

@@ -130,3 +130,7 @@ func (s *KafkaSourceStatus) MarkKeyTypeCorrect() {
 func (s *KafkaSourceStatus) MarkKeyTypeIncorrect(reason, messageFormat string, messageA ...interface{}) {
  KafkaSourceCondSet.Manage(s).MarkFalse(KafkaConditionKeyType, reason, messageFormat, messageA...)
 }
+
+func (s *KafkaSourceStatus) UpdateConsumerGroupStatus(status string) {
+ s.Claims = status
+}

pkg/apis/sources/v1beta1/kafka_types.go

@@ -113,6 +113,10 @@ type KafkaSourceStatus struct {
  // +optional
  Selector string `json:"selector,omitempty"`
 
+ // Claims consumed by this KafkaSource instance
+ // +optional
+ Claims string `json:"claims,omitempty"`
+
  // Implement Placeable.
  // +optional
  v1alpha1.Placeable `json:",inline"`

pkg/channel/consolidated/dispatcher/dispatcher_test.go

@@ -45,7 +45,7 @@ type mockKafkaConsumerFactory struct {
  createErr bool
 }
 
-func (c mockKafkaConsumerFactory) StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler consumer.KafkaConsumerHandler) (sarama.ConsumerGroup, error) {
+func (c mockKafkaConsumerFactory) StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler consumer.KafkaConsumerHandler, options ...consumer.SaramaConsumerHandlerOption) (sarama.ConsumerGroup, error) {
  if c.createErr {
  return nil, errors.New("error creating consumer")
  }

pkg/common/consumer/consumer_factory.go

@@ -27,7 +27,7 @@ var newConsumerGroup = sarama.NewConsumerGroup
 
 // Kafka consumer factory creates the ConsumerGroup and start consuming the specified topic
 type KafkaConsumerGroupFactory interface {
- StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler KafkaConsumerHandler) (sarama.ConsumerGroup, error)
+ StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler KafkaConsumerHandler, options ...SaramaConsumerHandlerOption) (sarama.ConsumerGroup, error)
 }
 
 type kafkaConsumerGroupFactoryImpl struct {
@@ -53,7 +53,7 @@ func (c *customConsumerGroup) Close() error {
 
 var _ sarama.ConsumerGroup = (*customConsumerGroup)(nil)
 
-func (c kafkaConsumerGroupFactoryImpl) StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler KafkaConsumerHandler) (sarama.ConsumerGroup, error) {
+func (c kafkaConsumerGroupFactoryImpl) StartConsumerGroup(groupID string, topics []string, logger *zap.SugaredLogger, handler KafkaConsumerHandler, options ...SaramaConsumerHandlerOption) (sarama.ConsumerGroup, error) {
  consumerGroup, err := newConsumerGroup(c.addrs, groupID, c.config)
  if err != nil {
  return nil, err
@@ -65,7 +65,7 @@ func (c kafkaConsumerGroupFactoryImpl) StartConsumerGroup(groupID string, topics
  go func() {
  defer close(errorCh)
  for {
- consumerHandler := NewConsumerHandler(logger, handler, errorCh)
+ consumerHandler := NewConsumerHandler(logger, handler, errorCh, options...)
 
  err := consumerGroup.Consume(ctx, topics, &consumerHandler)
  if err == sarama.ErrClosedConsumerGroup {

pkg/common/consumer/consumer_handler.go

@@ -32,29 +32,61 @@ type KafkaConsumerHandler interface {
  GetConsumerGroup() string
 }
 
+type SaramaConsumerLifecycleListener interface {
+ // Setup is invoked when the consumer is joining the session
+ Setup(sess sarama.ConsumerGroupSession)
+
+ // Cleanup is invoked when the consumer is leaving the session
+ Cleanup(sess sarama.ConsumerGroupSession)
+}
+
+type noopSaramaConsumerLifecycleListener struct{}
+
+func (n noopSaramaConsumerLifecycleListener) Setup(sess sarama.ConsumerGroupSession) {}
+
+func (n noopSaramaConsumerLifecycleListener) Cleanup(sess sarama.ConsumerGroupSession) {}
+
+func WithSaramaConsumerLifecycleListener(listener SaramaConsumerLifecycleListener) SaramaConsumerHandlerOption {
+ return func(handler *SaramaConsumerHandler) {
+ handler.lifecycleListener = listener
+ }
+}
+
 // ConsumerHandler implements sarama.ConsumerGroupHandler and provides some glue code to simplify message handling
 // You must implement KafkaConsumerHandler and create a new SaramaConsumerHandler with it
 type SaramaConsumerHandler struct {
  // The user message handler
  handler KafkaConsumerHandler
 
+ lifecycleListener SaramaConsumerLifecycleListener
+
  logger *zap.SugaredLogger
 
  // Errors channel
  errors chan error
 }
 
-func NewConsumerHandler(logger *zap.SugaredLogger, handler KafkaConsumerHandler, errorsCh chan error) SaramaConsumerHandler {
- return SaramaConsumerHandler{
- logger: logger,
- handler: handler,
- errors: errorsCh,
+type SaramaConsumerHandlerOption func(*SaramaConsumerHandler)
+
+func NewConsumerHandler(logger *zap.SugaredLogger, handler KafkaConsumerHandler, errorsCh chan error, options ...SaramaConsumerHandlerOption) SaramaConsumerHandler {
+ sch := SaramaConsumerHandler{
+ handler: handler,
+ lifecycleListener: noopSaramaConsumerLifecycleListener{},
+ logger: logger,
+ errors: errorsCh,
  }
+
+ for _, f := range options {
+ f(&sch)
+ }
+
+ return sch
 }
 
 // Setup is run at the beginning of a new session, before ConsumeClaim
-func (consumer *SaramaConsumerHandler) Setup(sarama.ConsumerGroupSession) error {
+func (consumer *SaramaConsumerHandler) Setup(session sarama.ConsumerGroupSession) error {
  consumer.logger.Info("setting up handler")
+ consumer.lifecycleListener.Setup(session)
  return nil
 }
 
@@ -68,6 +100,7 @@ func (consumer *SaramaConsumerHandler) Cleanup(session sarama.ConsumerGroupSessi
  consumer.handler.SetReady(p, false)
  }
  }
+ consumer.lifecycleListener.Cleanup(session)
  return nil
 }
 

pkg/source/adapter/adapter.go

@@ -26,6 +26,9 @@ import (
 
  "golang.org/x/time/rate"
 
+ ctrl "knative.dev/control-protocol/pkg"
+ ctrlnetwork "knative.dev/control-protocol/pkg/network"
+
  "github.com/Shopify/sarama"
  "go.opencensus.io/trace"
  "go.uber.org/zap"
@@ -38,6 +41,7 @@ import (
 
  "knative.dev/eventing-kafka/pkg/common/consumer"
  "knative.dev/eventing-kafka/pkg/source/client"
+ kafkasourcecontrol "knative.dev/eventing-kafka/pkg/source/control"
 )
 
 const (
@@ -59,7 +63,9 @@ func NewEnvConfig() adapter.EnvConfigAccessor {
 }
 
 type Adapter struct {
- config *AdapterConfig
+ config *AdapterConfig
+ controlServer *ctrlnetwork.ControlServer
+
  httpMessageSender *kncloudevents.HTTPMessageSender
  reporter pkgsource.StatsReporter
  logger *zap.SugaredLogger
@@ -68,6 +74,8 @@ type Adapter struct {
 }
 
 var _ adapter.MessageAdapter = (*Adapter)(nil)
+var _ consumer.KafkaConsumerHandler = (*Adapter)(nil)
+var _ consumer.SaramaConsumerLifecycleListener = (*Adapter)(nil)
 var _ adapter.MessageAdapterConstructor = NewAdapter
 
 func NewAdapter(ctx context.Context, processed adapter.EnvConfigAccessor, httpMessageSender *kncloudevents.HTTPMessageSender, reporter pkgsource.StatsReporter) adapter.MessageAdapter {
@@ -86,7 +94,7 @@ func (a *Adapter) GetConsumerGroup() string {
  return a.config.ConsumerGroup
 }
 
-func (a *Adapter) Start(ctx context.Context) error {
+func (a *Adapter) Start(ctx context.Context) (err error) {
  a.logger.Infow("Starting with config: ",
  zap.String("Topics", strings.Join(a.config.Topics, ",")),
  zap.String("ConsumerGroup", a.config.ConsumerGroup),
@@ -95,14 +103,27 @@ func (a *Adapter) Start(ctx context.Context) error {
  zap.String("Namespace", a.config.Namespace),
  )
 
+ // Init control service
+ a.controlServer, err = ctrlnetwork.StartInsecureControlServer(ctx)
+ if err != nil {
+ return err
+ }
+ a.controlServer.MessageHandler(a)
+
  // init consumer group
  addrs, config, err := client.NewConfigWithEnv(context.Background(), &a.config.KafkaEnvConfig)
  if err != nil {
  return fmt.Errorf("failed to create the config: %w", err)
  }
 
  consumerGroupFactory := consumer.NewConsumerGroupFactory(addrs, config)
- group, err := consumerGroupFactory.StartConsumerGroup(a.config.ConsumerGroup, a.config.Topics, a.logger, a)
+ group, err := consumerGroupFactory.StartConsumerGroup(
+ a.config.ConsumerGroup,
+ a.config.Topics,
+ a.logger,
+ a,
+ consumer.WithSaramaConsumerLifecycleListener(a),
+ )
  if err != nil {
  panic(err)
  }
@@ -177,3 +198,22 @@ func (a *Adapter) Handle(ctx context.Context, msg *sarama.ConsumerMessage) (bool
 func (a *Adapter) SetRateLimits(r rate.Limit, b int) {
  a.rateLimiter = rate.NewLimiter(r, b)
 }
+
+func (a *Adapter) HandleServiceMessage(ctx context.Context, message ctrl.ServiceMessage) {
+ // In this first PR, there is only the RA sending messages to control plane,
+ // there is no message the control plane should send to the RA
+ a.logger.Info("Received unexpected control message")
+ message.Ack()
+}
+
+func (a *Adapter) Setup(sess sarama.ConsumerGroupSession) {
+ if err := a.controlServer.SendAndWaitForAck(kafkasourcecontrol.NotifySetupClaimsOpCode, kafkasourcecontrol.Claims(sess.Claims())); err != nil {
+ a.logger.Warnf("Cannot send the claims update: %v", err)
+ }
+}
+
+func (a *Adapter) Cleanup(sess sarama.ConsumerGroupSession) {
+ if err := a.controlServer.SendAndWaitForAck(kafkasourcecontrol.NotifyCleanupClaimsOpCode, kafkasourcecontrol.Claims(sess.Claims())); err != nil {
+ a.logger.Warnf("Cannot send the claims update: %v", err)
+ }
+}

pkg/source/control/message.go

@@ -0,0 +1,91 @@
+package control
+
+import (
+ "encoding/json"
+ "fmt"
+ "strings"
+
+ "k8s.io/apimachinery/pkg/util/sets"
+
+ ctrl "knative.dev/control-protocol/pkg"
+)
+
+// This just contains the different opcodes
+const (
+ NotifySetupClaimsOpCode ctrl.OpCode = 1
+ NotifyCleanupClaimsOpCode ctrl.OpCode = 2
+)
+
+type Claims map[string][]int32
+
+func ClaimsParser(payload []byte) (interface{}, error) {
+ var claims Claims
+ err := (&claims).UnmarshalBinary(payload)
+ if err != nil {
+ return nil, err
+ }
+ return claims, nil
+}
+
+func ClaimsMerger(old interface{}, new interface{}) interface{} {
+ oldClaims := old.(Claims)
+ newClaims := new.(Claims)
+ result := oldClaims.copy()
+
+ for topic, partitions := range result {
+ if newPartitions, ok := newClaims[topic]; ok {
+ result[topic] = sets.NewInt32(partitions...).Insert(newPartitions...).List() // Merge partitions
+ delete(newClaims, topic)
+ }
+ }
+ for newTopic, newPartitions := range newClaims {
+ result[newTopic] = newPartitions
+ }
+
+ return result
+}
+
+func ClaimsDifference(old interface{}, new interface{}) interface{} {
+ oldClaims := old.(Claims)
+ cleanedClaims := new.(Claims)
+ result := oldClaims.copy()
+
+ for topic, partitions := range result {
+ if cleanedPartitions, ok := cleanedClaims[topic]; ok {
+ newSet := sets.NewInt32(partitions...).Delete(cleanedPartitions...).List()
+ if len(newSet) == 0 {
+ delete(result, topic)
+ } else {
+ result[topic] = newSet
+ }
+ }
+ }
+
+ return result
+}
+
+func (c Claims) String() string {
+ strs := make([]string, 0, len(c))
+ for topic, partitions := range c {
+ strs = append(strs, fmt.Sprintf("'%s': %v", topic, partitions))
+ }
+ return strings.Join(strs, ", ")
+}
+
+func (c Claims) MarshalBinary() (data []byte, err error) {
+ return json.Marshal(c)
+}
+
+func (c *Claims) UnmarshalBinary(data []byte) error {
+ return json.Unmarshal(data, c)
+}
+
+func (c Claims) copy() Claims {
+ res := make(Claims, len(c))
+
+ for k, v := range c {
+ res[k] = v
+ }
+
+ return res
+}