Skip to content

Commit 191e826

Browse files
vadorovskyvadorovsky
authored
Merge pull request #29 from deepfence/switch-kafka-lib
kafka: Replace confluentinc/confluent-kafka-go with segmentio/kafka-go
2 parents cbc3c1b + b9dc1f4 commit 191e826

File tree

5 files changed

+163
-75
lines changed

5 files changed

+163
-75
lines changed

‎go.mod‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@ require (
66
github.com/aws/aws-sdk-go-v2 v1.16.4
77
github.com/aws/aws-sdk-go-v2/config v1.15.3
88
github.com/aws/aws-sdk-go-v2/service/s3 v1.26.3
9-
github.com/confluentinc/confluent-kafka-go v1.8.2
109
github.com/foxcpp/go-mockdns v1.0.0
1110
github.com/google/gopacket v1.1.19
1211
github.com/google/uuid v1.3.0
1312
github.com/inhies/go-bytesize v0.0.0-20210819104631-275770b98743
14-
github.com/klauspost/compress v1.14.1
13+
github.com/klauspost/compress v1.14.2
14+
github.com/segmentio/kafka-go v0.4.32
1515
github.com/spf13/cobra v1.4.0
1616
gopkg.in/yaml.v3 v3.0.1
1717
)
@@ -33,6 +33,7 @@ require (
3333
github.com/inconshreveable/mousetrap v1.0.0 // indirect
3434
github.com/kr/pretty v0.1.0 // indirect
3535
github.com/miekg/dns v1.1.25 // indirect
36+
github.com/pierrec/lz4/v4 v4.1.14 // indirect
3637
github.com/spf13/pflag v1.0.5 // indirect
3738
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550 // indirect
3839
golang.org/x/net v0.0.0-20211209124913-491a49abca63 // indirect

‎go.sum‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
github.com/aws/aws-sdk-go-v2 v1.16.2 h1:fqlCk6Iy3bnCumtrLz9r3mJ/2gUT0pJ0wLFVIdWh+JA=
21
github.com/aws/aws-sdk-go-v2 v1.16.2/go.mod h1:ytwTPBG6fXTZLxxeeCCWj2/EMYp/xDUgX+OET6TLNNU=
32
github.com/aws/aws-sdk-go-v2 v1.16.4 h1:swQTEQUyJF/UkEA94/Ga55miiKFoXmm/Zd67XHgmjSg=
43
github.com/aws/aws-sdk-go-v2 v1.16.4/go.mod h1:ytwTPBG6fXTZLxxeeCCWj2/EMYp/xDUgX+OET6TLNNU=
@@ -32,10 +31,10 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.16.3 h1:cJGRyzCSVwZC7zZZ1xbx9m32UnrK
3231
github.com/aws/aws-sdk-go-v2/service/sts v1.16.3/go.mod h1:bfBj0iVmsUyUg4weDB4NxktD9rDGeKSVWnjTnwbx9b8=
3332
github.com/aws/smithy-go v1.11.2 h1:eG/N+CcUMAvsdffgMvjMKwfyDzIkjM6pfxMJ8Mzc6mE=
3433
github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM=
35-
github.com/confluentinc/confluent-kafka-go v1.8.2 h1:PBdbvYpyOdFLehj8j+9ba7FL4c4Moxn79gy9cYKxG5E=
36-
github.com/confluentinc/confluent-kafka-go v1.8.2/go.mod h1:u2zNLny2xq+5rWeTQjFHbDzzNuba4P1vo31r9r4uAdg=
3734
github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
3835
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
36+
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
37+
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
3938
github.com/foxcpp/go-mockdns v1.0.0 h1:7jBqxd3WDWwi/6WhDvacvH1XsN3rOLXyHM1uhvIx6FI=
4039
github.com/foxcpp/go-mockdns v1.0.0/go.mod h1:lgRN6+KxQBawyIghpnl5CezHFGS9VLzvtVlwxvzXTQ4=
4140
github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
@@ -50,23 +49,35 @@ github.com/inhies/go-bytesize v0.0.0-20210819104631-275770b98743 h1:X3Xxno5Ji8id
5049
github.com/inhies/go-bytesize v0.0.0-20210819104631-275770b98743/go.mod h1:KrtyD5PFj++GKkFS/7/RRrfnRhAMGQwy75GLCHWrCNs=
5150
github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
5251
github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
53-
github.com/klauspost/compress v1.14.1 h1:hLQYb23E8/fO+1u53d02A97a8UnsddcvYzq4ERRU4ds=
54-
github.com/klauspost/compress v1.14.1/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
52+
github.com/klauspost/compress v1.14.2 h1:S0OHlFk/Gbon/yauFJ4FfJJF5V0fc5HbBTJazi28pRw=
53+
github.com/klauspost/compress v1.14.2/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
5554
github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
5655
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
5756
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
5857
github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
5958
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
6059
github.com/miekg/dns v1.1.25 h1:dFwPR6SfLtrSwgDcIq2bcU/gVutB4sNApq2HBdqcakg=
6160
github.com/miekg/dns v1.1.25/go.mod h1:bPDLeHnStXmXAq1m/Ch/hvfNHr14JKNPMBo3VZKjuso=
61+
github.com/pierrec/lz4/v4 v4.1.14 h1:+fL8AQEZtz/ijeNnpduH0bROTu0O3NZAlPjQxGn8LwE=
62+
github.com/pierrec/lz4/v4 v4.1.14/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
63+
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
6264
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
6365
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
66+
github.com/segmentio/kafka-go v0.4.32 h1:Ohr+9E+kDv/Ld2UPJN9hnKZRd2qgiqCmI8v2e1qlfLM=
67+
github.com/segmentio/kafka-go v0.4.32/go.mod h1:JAPPIiY3MQIwVHj64CWOP0LsFFfQ7H0w69kuoxnMIS0=
6468
github.com/spf13/cobra v1.4.0 h1:y+wJpx64xcgO1V+RcnwW0LEHxTKRi2ZDPSBjWnrg88Q=
6569
github.com/spf13/cobra v1.4.0/go.mod h1:Wo4iy3BUC+X2Fybo0PDqwJIv3dNRiZLHQymsfxlB84g=
6670
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
6771
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
6872
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
73+
github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
74+
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
75+
github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c h1:u40Z8hqBAAQyv+vATcGgV0YCnDjqSL7/q/JyPhhJSPk=
76+
github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c/go.mod h1:lB8K/P019DLNhemzwFU4jHLhdvlE6uDZjXFejJXr49I=
77+
github.com/xdg/stringprep v1.0.0 h1:d9X0esnoa3dFsV0FG35rAT0RIhYFlPq7MiP+DW89La0=
78+
github.com/xdg/stringprep v1.0.0/go.mod h1:Jhud4/sHMO4oL310DaZAKk9ZaJ08SJfe+sJh0HrGL1Y=
6979
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
80+
golang.org/x/crypto v0.0.0-20190506204251-e1dfcc566284/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
7081
golang.org/x/crypto v0.0.0-20190923035154-9ee001bba392/go.mod h1:/lpIB1dKB+9EgE3H3cr1v9wB50oz8l4C4h62xy7jSTY=
7182
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550 h1:ObdrDkeb4kJdCP557AjRjq69pTHfNouLtWZG7j9rPN8=
7283
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
@@ -90,6 +101,7 @@ golang.org/x/sys v0.0.0-20211210111614-af8b64212486/go.mod h1:oPkhp1MJrh7nUepCBc
90101
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
91102
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
92103
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
104+
golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
93105
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
94106
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
95107
golang.org/x/tools v0.0.0-20190907020128-2ca718005c18/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
@@ -103,7 +115,7 @@ gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogR
103115
gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
104116
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
105117
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
106-
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
107-
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
118+
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
119+
gopkg.in/yaml.v3 v3.0.0-20220512140231-539c8e751b99/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
108120
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
109121
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

‎pkg/config/config.go‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@ package config
22

33
import (
44
"fmt"
5-
"github.com/aws/aws-sdk-go-v2/service/s3/types"
6-
"github.com/inhies/go-bytesize"
75
"io/ioutil"
6+
"strings"
87
"time"
98

9+
"github.com/aws/aws-sdk-go-v2/service/s3/types"
10+
"github.com/inhies/go-bytesize"
11+
1012
"github.com/klauspost/compress/s2"
1113
"gopkg.in/yaml.v3"
1214
)
@@ -53,12 +55,13 @@ type S3PluginConfig struct {
5355
}
5456

5557
type KafkaPluginConfig struct {
56-
Brokers string
58+
Brokers []string
5759
ClientId string `yaml:"clientId,omitempty"`
5860
Topic string `yaml:"topic,omitempty"`
5961
MessageSize *bytesize.ByteSize `yaml:"messageSize,omitempty"`
6062
Acks string `yaml:"acks,omitempty"`
6163
FileSize *bytesize.ByteSize `yaml:"fileSize,omitempty"`
64+
Timeout time.Duration `yaml:"timeout,omitempty"`
6265
}
6366

6467
type PluginsConfig struct {
@@ -83,11 +86,12 @@ type S3OutputRawConfig struct {
8386

8487
type KafkaOutputRawConfig struct {
8588
Brokers string
86-
ClientId *string `yaml:"clientId,omitempty"`
87-
Topic *string `yaml:"topic,omitempty"`
88-
MessageSize *string `yaml:"messageSize,omitempty"`
89-
Acks *string `yaml:"acks,omitempty"`
90-
FileSize *string `yaml:"fileSize,omitempty"`
89+
ClientId *string `yaml:"clientId,omitempty"`
90+
Topic *string `yaml:"topic,omitempty"`
91+
MessageSize *string `yaml:"messageSize,omitempty"`
92+
Acks *string `yaml:"acks,omitempty"`
93+
FileSize *string `yaml:"fileSize,omitempty"`
94+
Timeout time.Duration `yaml:"timeout,omitempty"`
9195
}
9296

9397
type PluginsRawConfig struct {
@@ -289,12 +293,13 @@ func populateKafkaConfig(rawConfig RawConfig) (*KafkaPluginConfig, error) {
289293
}
290294

291295
return &KafkaPluginConfig{
292-
Brokers: rawConfig.Output.Plugins.Kafka.Brokers,
296+
Brokers: strings.Split(rawConfig.Output.Plugins.Kafka.Brokers, ","),
293297
ClientId: clientId,
294298
Topic: topic,
295299
MessageSize: messageSize,
296300
Acks: acks,
297301
FileSize: fileSize,
302+
Timeout: rawConfig.Output.Plugins.Kafka.Timeout,
298303
}, nil
299304
}
300305

‎pkg/plugins/kafka/kafka.go‎

Lines changed: 38 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,18 @@ package kafka
22

33
import (
44
"context"
5-
"fmt"
6-
"github.com/confluentinc/confluent-kafka-go/kafka"
5+
"log"
6+
"time"
7+
78
"github.com/deepfence/PacketStreamer/pkg/config"
89
"github.com/deepfence/PacketStreamer/pkg/file"
910
"github.com/google/uuid"
10-
"log"
11+
kafka "github.com/segmentio/kafka-go"
1112
)
1213

13-
type KafkaProducer interface {
14-
Produce(msg *kafka.Message, deliveryChan chan kafka.Event) error
15-
Close()
14+
type KafkaWriter interface {
15+
WriteMessages(ctx context.Context, msgs ...kafka.Message) error
16+
Close() error
1617
}
1718

1819
type File struct {
@@ -25,8 +26,19 @@ func (f *File) newBuffer(size int) {
2526
f.Buffer = make([]byte, 0, size)
2627
}
2728

29+
type IdGenerator interface {
30+
Generate() string
31+
}
32+
33+
type FileIdGenerator struct{}
34+
35+
func (g *FileIdGenerator) Generate() string {
36+
return uuid.New().String()
37+
}
38+
2839
type Plugin struct {
29-
Producer KafkaProducer
40+
Writer KafkaWriter
41+
IdGenerator IdGenerator
3042
Topic string
3143
MessageSize int
3244
CloseChan chan bool
@@ -35,18 +47,20 @@ type Plugin struct {
3547
}
3648

3749
func NewPlugin(config *config.KafkaPluginConfig) (*Plugin, error) {
38-
producer, err := kafka.NewProducer(&kafka.ConfigMap{
39-
"bootstrap.servers": config.Brokers,
40-
"client.id": config.ClientId,
41-
"acks": config.Acks,
42-
})
43-
44-
if err != nil {
45-
return nil, fmt.Errorf("error creating kafka producer, %w", err)
50+
dialer := &kafka.Dialer{
51+
Timeout: 10 * time.Second,
52+
DualStack: true,
53+
TLS: nil,
4654
}
55+
writer := kafka.NewWriter(kafka.WriterConfig{
56+
Brokers: config.Brokers,
57+
Topic: config.Topic,
58+
Balancer: &kafka.Hash{},
59+
Dialer: dialer,
60+
})
4761

4862
return &Plugin{
49-
Producer: producer,
63+
Writer: writer,
5064
Topic: config.Topic,
5165
MessageSize: int(*config.MessageSize),
5266
FileSize: uint64(*config.FileSize),
@@ -67,8 +81,8 @@ func (p *Plugin) newFile(id string, messageSize int) {
6781
func (p *Plugin) Start(ctx context.Context) chan<- string {
6882
inputChan := make(chan string)
6983
go func() {
70-
defer p.Producer.Close()
71-
p.newFile(generateFileId(), p.MessageSize)
84+
defer p.Writer.Close()
85+
p.newFile(p.IdGenerator.Generate(), p.MessageSize)
7286

7387
for {
7488
select {
@@ -103,7 +117,7 @@ func (p *Plugin) Start(ctx context.Context) chan<- string {
103117
}
104118

105119
if p.CurrentFile.Sent >= p.FileSize {
106-
p.newFile(generateFileId(), p.MessageSize)
120+
p.newFile(p.IdGenerator.Generate(), p.MessageSize)
107121
} else {
108122
p.CurrentFile.newBuffer(p.MessageSize)
109123
}
@@ -132,17 +146,13 @@ func (p *Plugin) cleanup() {
132146
}
133147

134148
func (p *Plugin) flush() error {
135-
err := p.Producer.Produce(&kafka.Message{
136-
TopicPartition: kafka.TopicPartition{Topic: &p.Topic, Partition: kafka.PartitionAny},
137-
Value: p.CurrentFile.Buffer,
138-
Key: []byte(p.CurrentFile.Id),
139-
}, nil)
149+
err := p.Writer.WriteMessages(context.Background(), kafka.Message{
150+
Topic: p.Topic,
151+
Key: []byte(p.CurrentFile.Id),
152+
Value: p.CurrentFile.Buffer,
153+
})
140154

141155
p.CurrentFile.Sent += uint64(len(p.CurrentFile.Buffer))
142156

143157
return err
144158
}
145-
146-
func generateFileId() string {
147-
return uuid.New().String()
148-
}

0 commit comments

Comments
 (0)