80 lines
2.0 KiB
Go
80 lines
2.0 KiB
Go
package provider
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/Shopify/sarama"
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
// ProducerProvider pool of producers that ensure transactional-id is unique.
|
|
type ProducerProvider struct {
|
|
transactionIdGenerator int32
|
|
|
|
producersLock sync.Mutex
|
|
producers []sarama.SyncProducer
|
|
|
|
producerProvider func() sarama.SyncProducer
|
|
}
|
|
|
|
func NewProducerProvider(brokers []string, producerConfigurationProvider func() *sarama.Config) Provider[sarama.SyncProducer] {
|
|
provider := &ProducerProvider{}
|
|
provider.producerProvider = func() sarama.SyncProducer {
|
|
configuration := producerConfigurationProvider()
|
|
suffix := provider.transactionIdGenerator
|
|
// Append transactionIdGenerator to current configuration.Producer.Transaction.ID to ensure transaction-id uniqueness.
|
|
if configuration.Producer.Transaction.ID != "" {
|
|
provider.transactionIdGenerator++
|
|
configuration.Producer.Transaction.ID = configuration.Producer.Transaction.ID + "-" + fmt.Sprint(suffix)
|
|
}
|
|
producer, err := sarama.NewSyncProducer(brokers, configuration)
|
|
if err != nil {
|
|
log.Fatalln(err)
|
|
}
|
|
return producer
|
|
}
|
|
return provider
|
|
}
|
|
|
|
func (p *ProducerProvider) Borrow() (producer sarama.SyncProducer) {
|
|
p.producersLock.Lock()
|
|
defer p.producersLock.Unlock()
|
|
|
|
if len(p.producers) == 0 {
|
|
for {
|
|
producer = p.producerProvider()
|
|
if producer != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
index := len(p.producers) - 1
|
|
producer = p.producers[index]
|
|
p.producers = p.producers[:index]
|
|
return
|
|
}
|
|
|
|
func (p *ProducerProvider) Release(producer sarama.SyncProducer) {
|
|
p.producersLock.Lock()
|
|
defer p.producersLock.Unlock()
|
|
|
|
// If released producer is erroneous close it and don't return it to the producer pool.
|
|
if producer.TxnStatus()&sarama.ProducerTxnFlagInError != 0 {
|
|
// Try to close it
|
|
_ = producer.Close()
|
|
return
|
|
}
|
|
p.producers = append(p.producers, producer)
|
|
}
|
|
|
|
func (p *ProducerProvider) Clear() {
|
|
p.producersLock.Lock()
|
|
defer p.producersLock.Unlock()
|
|
|
|
for _, producer := range p.producers {
|
|
_ = producer.Close()
|
|
}
|
|
p.producers = p.producers[:0]
|
|
}
|