qmp-testing-suite/golang-producer-consumer/vendor/gitlab-app.eng.qops.net/golang/qmp/client.go

// Package qpcl creates producers and consumers compliant with
// the QPCL specification.
package qpcl

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	uuid "github.com/satori/go.uuid"
	"gitlab-app.eng.qops.net/golang/metrics"
	"gitlab-app.eng.qops.net/golang/qmp/internal/config"
	qc "gitlab-app.eng.qops.net/golang/qmp/internal/qconsumer"
	qp "gitlab-app.eng.qops.net/golang/qmp/internal/qproducer"
	"gitlab-app.eng.qops.net/golang/qmp/internal/types"
	qsl "gitlab-app.eng.qops.net/golang/qmp/qsl"
	"gitlab-app.eng.qops.net/golang/qmp/qsl/beacon"
	"gitlab-app.eng.qops.net/golang/qmp/qsl/logger"
)

// QClientFactory creates producers and consumers compliant with
// the QPCL specification.
type QClientFactory interface {
	NewProducer(types.QPCType, string) (Producer, error)
	NewConsumer(types.QPCType, string, string, qsl.Reader, int) (Consumer, error)
	update(types.QPCType, map[string]string) error
	buildClientID(types.QPCType) (string, error)
}

// Producer puts QMessages into the configured topic and Kafka
// broker.
type Producer interface {
	Close() error
	IsClosed() bool
	GetType() types.QPCType
	GetTopic() string
	Put(qsl.Message, func(error, qsl.Message)) error
	GetClientID() string
}

// Consumer spawns goroutines to handle messages received from
// Kafka.
type Consumer interface {
	Start(func(qsl.Message)) error
	Stop() error
	Join() error
	IsStopped() bool
	GetTopic() string
	GetType() types.QPCType
	GetClientID() string
}

type updatable interface {
	Update(map[string]string) error
	GetClientID() string
	IsClosed() bool
}

// QKafkaConfig creates an interface for
// accessing the Kafka-configured settings for all
// producers or all consumers.
type QKafkaConfig interface {
	GetProperties(types.QPCType) (map[string]string, error)
	SetTypeConfiguration(types.QPCType, map[string]string) error
	SetTypeValue(types.QPCType, string, string) error

	GetOverride(string, string) string
	GetOverrideInt(string, int) (int, error)
	ApplyEnvironmentOverrides(string, map[string]string)

	HasType(types.QPCType) error
}

type client struct {
	appName          string
	producerConfig   QKafkaConfig
	consumerConfig   QKafkaConfig
	topicClassConfig QKafkaConfig
	reporter         metrics.Reporter
	clientIDs        map[string]bool
	configUpdates    Consumer

	manufactured map[types.QPCType][]updatable
	creates      int32
	updateLock   *sync.RWMutex
	clientIDLock *sync.Mutex
}

const (
	libraryName    = "qmp-go-client"
	libraryVersion = "unknown"
)

// instance is the singleton QClientFactory.
var instance QClientFactory
var once sync.Once

// ErrUnregisteredApplication indicates an attempt to NewConsumer, NewProducer, or Update without
// calling RegisterApplication.
var ErrUnregisteredApplication = errors.New("cannot use qpcl without calling RegisterApplication first")

// ErrAlreadyRegistered indicates an additional attempt to RegisterApplication.
var ErrAlreadyRegistered = errors.New("qpcl.RegisterApplication should only be called ONCE")

// RegisterApplication initializes the singleton QClientFactory.
func RegisterApplication(appName string) error {
	if instance != nil {
		return ErrAlreadyRegistered
	}
	v, err := newQClientFactory(appName)
	if err == nil {
		once.Do(func() {
			instance = v
			beacon.Register(appName, libraryName, libraryVersion, beacon.MetricsPrefix)
		})
	}
	return err
}

func newQClientFactory(appName string) (*client, error) {
	pConfig, err := config.NewQProducerConfig(getEnvironment())
	if err != nil {
		return nil, err
	}
	cConfig, err := config.NewQConsumerConfig(getEnvironment())
	if err != nil {
		return nil, err
	}
	tcConfig, err := config.NewTopicClassConfig(getEnvironment())
	if err != nil {
		return nil, err
	}
	reporter, err := metrics.NewReporter(beacon.MetricsPrefix)
	if err != nil {
		return nil, err
	}
	manufactured := make(map[types.QPCType][]updatable)
	for i := range types.GetQCTypes() {
		manufactured[i] = []updatable{}
	}
	for i := range types.GetQPTypes() {
		manufactured[i] = []updatable{}
	}
	c := &client{
		appName:          appName,
		producerConfig:   pConfig,
		consumerConfig:   cConfig,
		topicClassConfig: tcConfig,
		reporter:         reporter,
		manufactured:     manufactured,
		creates:          int32(0),
		clientIDs:        make(map[string]bool),
		updateLock:       &sync.RWMutex{},
		clientIDLock:     &sync.Mutex{},
	}
	return c, nil
}

// NewProducer returns a producer ready to send messages to Kafka
// with Put(qsl.QMessage, func(error, qsl.QMessage)).
func NewProducer(producerType types.QPCType, topicName string) (Producer, error) {
	if instance == nil {
		return nil, ErrUnregisteredApplication
	}
	return instance.NewProducer(producerType, topicName)
}

// client.NewProducer returns a producer ready to send messages
// to Kafka with Put(qsl.QMessage, handler(error, qsl.QMessage)).
func (c *client) NewProducer(producerType types.QPCType, topicName string) (Producer, error) {
	c.updateLock.Lock()
	defer c.updateLock.Unlock()
	pconfig, err := c.getConfig(producerType)
	if err != nil {
		return nil, fmt.Errorf("unknown producer type %d", int(producerType))
	}
	qproducer, err := qp.NewProducer(producerType, topicName, pconfig, c.reporter)
	if err != nil {
		return nil, err
	}
	err = c.incCounterQPCNewd(c.reporter, producerType, topicName)
	if err != nil {
		return nil, err
	}
	c.manufactured[producerType] = append(c.manufactured[producerType], qproducer)
	return qproducer, nil
}

func (c *client) getConfig(producerConsumerType types.QPCType) (map[string]string, error) {
	pcconfig, err := c.producerConfig.GetProperties(producerConsumerType)
	if err != nil {
		pcconfig, err = c.consumerConfig.GetProperties(producerConsumerType)
	}
	if err != nil {
		return nil, err
	}
	clientID, err := c.buildClientID(producerConsumerType)
	if err != nil {
		return nil, err
	}
	pcconfig[config.ClientID] = clientID
	return pcconfig, nil
}

// NewConsumer returns a consumer ready to receive messages from Kafka
// with Start(func(qsl.QMessage)).
func NewConsumer(consumerType types.QPCType, topicName, consumerGroup string, reader qsl.Reader, threads int) (Consumer, error) {
	if instance == nil {
		return nil, ErrUnregisteredApplication
	}
	return instance.NewConsumer(consumerType, topicName, consumerGroup, reader, threads)
}

// client.NewConsumer creates a consumer ready to process messages
// from Kafka with Start(handler(qsl.QMessage)).
func (c *client) NewConsumer(consumerType types.QPCType, topicName, consumerGroup string, reader qsl.Reader, threads int) (Consumer, error) {
	c.updateLock.Lock()
	defer c.updateLock.Unlock()
	cconfig, err := c.getConfig(consumerType)
	if err != nil {
		return nil, fmt.Errorf("unknown consumer type %d", int(consumerType))
	}
	if _, ok := cconfig[config.QMPThreads]; !ok {
		cconfig[config.QMPThreads] = strconv.Itoa(threads)
	}
	qconsumer, err := qc.NewConsumer(topicName, consumerGroup, reader, cconfig, c.reporter, consumerType)
	if err != nil {
		return nil, err
	}
	err = c.incCounterQPCNewd(c.reporter, consumerType, topicName)
	if err != nil {
		return nil, err
	}
	c.manufactured[consumerType] = append(c.manufactured[consumerType], qconsumer)
	return qconsumer, nil
}

func (c *client) countNew() {
	producersConsumersCreated := atomic.AddInt32(&c.creates, 1)
	if producersConsumersCreated > int32(5) {
		c.clean()
		atomic.StoreInt32(&c.creates, 0)
	}
}

// Update iterates over all returned producers and consumers and forgets them
// if closed or updates them with the given newConfiguration if the producer or
// consumer's type matches the given type.
func Update(consumerProducerType types.QPCType, newConfiguration map[string]string) error {
	if instance == nil {
		return ErrUnregisteredApplication
	}
	return instance.update(consumerProducerType, newConfiguration)
}

// Update iterates over all returned producers and consumers and forgets them
// if closed or updates them if the type matches.
func (c *client) update(consumerProducerType types.QPCType, newConfiguration map[string]string) error {
	c.updateLock.Lock()
	defer c.updateLock.Unlock()
	c.clean()
	for i := range c.manufactured {
		if i == consumerProducerType {
			var qkc QKafkaConfig
			if c.producerConfig.HasType(consumerProducerType) == nil {
				qkc = c.producerConfig
			} else if c.consumerConfig.HasType(consumerProducerType) == nil {
				qkc = c.consumerConfig
			} else {
				break
			}
			return c.updateGivenType(i, newConfiguration, c.manufactured[i], qkc)
		}
	}
	return fmt.Errorf("cannot update configuration for unknown consumer or producer type %d", consumerProducerType)
}

func (c *client) updateGivenType(consProdType types.QPCType, newConf map[string]string, batch []updatable, consProdConf QKafkaConfig) error {
	previousConfig, err := consProdConf.GetProperties(consProdType)
	if err != nil {
		return err
	}
	err = consProdConf.SetTypeConfiguration(consProdType, newConf)
	if err != nil {
		consProdConf.SetTypeConfiguration(consProdType, previousConfig)
		return err
	}
	newConf, err = consProdConf.GetProperties(consProdType)
	if err != nil {
		consProdConf.SetTypeConfiguration(consProdType, previousConfig)
		return err
	}
	l := len(batch)
	for i := 0; i < l; i++ {
		err = batch[i].Update(newConf)
		if err != nil && err != qc.ErrUpdateClosed && err != qp.ErrUpdateClosed {
			consProdConf.SetTypeConfiguration(consProdType, previousConfig)
			return err
		}
	}
	return nil
}

func (c *client) clean() {
	for k := range c.manufactured {
		l := len(c.manufactured[k])
		for i := 0; i < l; i++ {
			if c.manufactured[k][i].IsClosed() {
				c.makeClientIDAvailable(c.manufactured[k][i].GetClientID())
				l--
				c.manufactured[k][i] = c.manufactured[k][l]
				c.manufactured[k] = c.manufactured[k][:l]
				i--
			}
		}
	}
}

func (c *client) buildClientID(producerConsumerType types.QPCType) (string, error) {
	typeString, err := producerConsumerType.String()
	if err != nil {
		return "", err
	}
	var validID bool
	var clientID string
	for !validID {
		uuid := uuid.NewV4().String()
		clientID = fmt.Sprintf(
			"qmp-%s-%s-%s",
			typeString,
			GetLibraryVersion(),
			uuid,
		)
		if c.clientIDAvailable(clientID) == nil {
			break
		}
	}
	c.clientIDs[clientID] = true
	return clientID, nil
}

func getEnvironment() map[string]string {
	m := make(map[string]string)
	env := os.Environ()
	for i := range env {
		env[i] = strings.TrimSpace(env[i])
		tokens := strings.SplitN(env[i], "=", 2)
		if len(tokens) < 2 {
			continue
		}
		m[tokens[0]] = tokens[1]
	}
	return m
}

// GetLibraryVersion returns the string of the qmp go library's version.
func GetLibraryVersion() string {
	return libraryVersion
}

// IsRegistered returns whether the QPCL client factory singleton has been
// registered or not.
func IsRegistered() bool {
	return instance != nil
}

// GetLibraryName returns the string of the qmp go library's name.
func GetLibraryName() string {
	return libraryName
}

func (c *client) clientIDAvailable(clientID string) error {
	c.clientIDLock.Lock()
	defer c.clientIDLock.Unlock()
	if _, ok := c.clientIDs[clientID]; ok {
		return fmt.Errorf("client ID %q exists", clientID)
	}
	return nil
}

func (c *client) makeClientIDAvailable(clientID string) {
	c.clientIDLock.Lock()
	defer c.clientIDLock.Unlock()
	if _, ok := c.clientIDs[clientID]; ok {
		delete(c.clientIDs, clientID)
	}
}

type updateItem struct {
	name  string
	typed types.QPCType
}

func (c *client) configurationUpdate(m qsl.Message) {
	payload, err := m.GetPayload()
	if err != nil {
		return
	}
	toUpdate := []updateItem{}
	// 2do topic class configuration... something.
	for k, v := range types.GetQCTypes() {
		toUpdate = append(toUpdate, updateItem{
			name:  strings.ToLower(v) + "_consumer",
			typed: k,
		})
	}
	for k, v := range types.GetQPTypes() {
		toUpdate = append(toUpdate, updateItem{
			name:  strings.ToLower(v) + "_producer",
			typed: k,
		})
	}
	for i := range toUpdate {
		m, err := config.MapAtMapKey(payload, toUpdate[i].name)
		if err != nil {
			logger.Warn(err)
			continue
		}
		confMap := map[string]string{}
		raisedMap := raiseMap(m)
		for k, v := range raisedMap {
			if v != nil {
				confMap[k] = fmt.Sprintf("%v", v)
			}
		}
		for k := range confMap {
			confMap[strings.Replace(k, "_", ".", -1)] = confMap[k]
			delete(confMap, strings.Replace(k, ".", "_", -1))
		}
		c.update(toUpdate[i].typed, confMap)
	}
}

func raiseMap(m map[string]interface{}) map[string]interface{} {
	out := make(map[string]interface{})
	for k, v := range m {
		switch v.(type) {
		case map[string]interface{}:
			replace := raiseMap(v.(map[string]interface{}))
			for replaceK := range replace {
				out[replaceK] = replace[replaceK]
			}
		case []interface{}:
			for _, arrItem := range v.([]interface{}) {
				if arrMap, ok := arrItem.(map[string]interface{}); ok {
					replace := raiseMap(arrMap)
					for replaceK := range replace {
						out[replaceK] = replace[replaceK]
					}
				}
			}
		default:
			out[k] = v
		}
	}
	return out
}

func (c *client) initStateTopicConsumer() error {
	c.updateLock.Lock()
	defer c.updateLock.Unlock()
	cconfig, err := c.getConfig(types.CRecordState)
	if err != nil {
		return err
	}
	cconfig[config.QMPThreads] = "1"
	qconsumer, err := qc.NewStateTopicConsumer(config.GroupStateTopic, cconfig, c.reporter)
	if err != nil {
		return err
	}
	err = c.incCounterQPCNewd(c.reporter, types.CRecordState, qconsumer.GetTopic())
	if err != nil {
		return err
	}
	err = qconsumer.Start(c.configurationUpdate)
	if err != nil {
		qconsumer.Stop()
		return err
	}
	c.manufactured[types.CRecordState] = append(c.manufactured[types.CRecordState], qconsumer)
	c.configUpdates = qconsumer
	return nil
}

func (c *client) incCounterQPCNewd(reporter metrics.Reporter, typed types.QPCType, topic string) error {
	c.countNew()
	typeString, err := typed.String()
	if err != nil {
		return err
	}
	var producerOrConsumer string
	for k := range types.GetQPTypes() {
		if k == typed {
			producerOrConsumer = "producer"
			break
		}
	}
	for k := range types.GetQCTypes() {
		if k == typed {
			producerOrConsumer = "consumer"
			break
		}
	}
	reporter.IncCounter(fmt.Sprintf("factory.%s.created", producerOrConsumer),
		metrics.Tag("type", strings.ToLower(typeString)),
		metrics.Tag("topic", topic),
	)
	return nil
}