kafka-go publisher and subsriber

Author: wuriyanto48

cmd/kafka-cli/main.go

@@ -30,14 +30,14 @@ func main() {
 	}
 
 	ctx := context.Background()
-	publisher, err := kafka.NewPublisher(args.Brokers...)
+	publisher, err := kafka.NewKafkaGoPublisher(args.Topic, args.Brokers...)
 	if err != nil {
 		fmt.Println("error : ", err)
 
 		os.Exit(1)
 	}
 
-	subsriber, err := kafka.NewSubscriber(args.Brokers...)
+	subsriber, err := kafka.NewKafkaGoSubscriber(args.Topic, args.Brokers...)
 	if err != nil {
 		fmt.Println("error : ", err)
 

kafka.go

@@ -22,7 +22,7 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	switch command {
 	case PublishCommand:
-		return r.Publisher.Publish(r.Argument.Topic, r.Argument.Message)
+		return r.Publisher.Publish(ctx, r.Argument.Topic, r.Argument.Message)
 	case SubscribeCommand:
 		return r.Subscriber.Subscribe(ctx, r.Argument.Topic)
 	}

kafka_go_publisher.go

@@ -0,0 +1,46 @@
+package kafka
+
+import (
+	"context"
+	"time"
+
+	ka "github.com/segmentio/kafka-go"
+	"github.com/segmentio/kafka-go/snappy"
+)
+
+//KafkaGoPublisherImpl struct
+type KafkaGoPublisherImpl struct {
+	writer *ka.Writer
+}
+
+//NewKafkaGoPublisherImpl constructor of KafkaGoPublisherImpl
+func NewKafkaGoPublisher(topic string, addresses ...string) (*KafkaGoPublisherImpl, error) {
+	config := ka.WriterConfig{
+		Brokers:          addresses,
+		Topic:            topic,
+		Balancer:         &ka.LeastBytes{},
+		CompressionCodec: snappy.NewCompressionCodec(),
+		BatchTimeout:     5 * time.Millisecond,
+		BatchBytes:       1000000,
+	}
+
+	writer := ka.NewWriter(config)
+	return &KafkaGoPublisherImpl{writer}, nil
+}
+
+//Publish function
+func (publisher *KafkaGoPublisherImpl) Publish(ctx context.Context, topic string, message []byte) error {
+	defer func() { publisher.writer.Close() }()
+
+	msg := ka.Message{
+		Topic: topic,
+		Value: message,
+	}
+
+	err := publisher.writer.WriteMessages(ctx, msg)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

kafka_go_subscriber.go

@@ -0,0 +1,88 @@
+package kafka
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/segmentio/kafka-go"
+	ka "github.com/segmentio/kafka-go"
+	_ "github.com/segmentio/kafka-go/snappy"
+)
+
+//KafkaGoSubscriberImpl struct
+type KafkaGoSubscriberImpl struct {
+	reader *ka.Reader
+}
+
+//NewSubscriber constructor of KafkaGoSubscriberImpl
+func NewKafkaGoSubscriber(topic string, addresses ...string) (*KafkaGoSubscriberImpl, error) {
+	config := kafka.ReaderConfig{
+		Brokers:  addresses,
+		Topic:    topic,
+		GroupID:  "kafka-cli-group",
+		MinBytes: 10e3, // 10KB
+		MaxBytes: 10e6, // 10MB
+		// CommitInterval: time.Second, // flushes commits to Kafka every second
+	}
+
+	reader := ka.NewReader(config)
+	return &KafkaGoSubscriberImpl{reader}, nil
+}
+
+//Subscribe function
+func (s *KafkaGoSubscriberImpl) Subscribe(ctx context.Context, topics ...string) error {
+	signals := make(chan os.Signal, 1)
+	done := make(chan bool, 1)
+
+	// wait until get signal from OS
+	signal.Notify(signals, os.Interrupt, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)
+
+	go waitOSNotify(signals, done)
+
+	go func() {
+		for {
+			message, err := s.reader.ReadMessage(ctx)
+			if err != nil {
+				if err == io.EOF {
+					fmt.Println("read message end ")
+					break
+				} else {
+					fmt.Println("read message error ", err)
+				}
+			}
+
+			fmt.Println()
+			fmt.Printf("\nMessage = %s\nTopic = %s", string(message.Value), message.Topic)
+			//s.reader.CommitMessages(ctx, message)
+		}
+	}()
+
+	<-done
+
+	// close subscriber after exit
+	err := s.close()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func waitOSNotify(kill chan os.Signal, done chan bool) {
+	select {
+	case <-kill:
+		fmt.Println("process interrupted")
+		done <- true
+		return
+	}
+
+}
+
+func (s *KafkaGoSubscriberImpl) close() error {
+	fmt.Println("closing subscriber")
+	return s.reader.Close()
+}