github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/broker/adapter/streams/output.go (about)

     1  // Copyright (c) 2021, R.I. Pienaar and the Choria Project contributors
     2  //
     3  // SPDX-License-Identifier: Apache-2.0
     4  
     5  package streams
     6  
     7  import (
     8  	"context"
     9  	"encoding/json"
    10  	"fmt"
    11  	"strconv"
    12  	"strings"
    13  	"sync"
    14  
    15  	"github.com/choria-io/go-choria/broker/adapter/ingest"
    16  	"github.com/choria-io/go-choria/broker/adapter/stats"
    17  	"github.com/choria-io/go-choria/broker/adapter/transformer"
    18  	"github.com/choria-io/go-choria/inter"
    19  	"github.com/choria-io/go-choria/internal/util"
    20  	"github.com/choria-io/go-choria/srvcache"
    21  	"github.com/prometheus/client_golang/prometheus"
    22  	"github.com/sirupsen/logrus"
    23  )
    24  
    25  type stream struct {
    26  	servers     func() (srvcache.Servers, error)
    27  	clientID    string
    28  	topic       string
    29  	conn        inter.Connector
    30  	log         *logrus.Entry
    31  	name        string
    32  	adapterName string
    33  
    34  	work chan ingest.Adaptable
    35  }
    36  
    37  func newStream(name string, work chan ingest.Adaptable, logger *logrus.Entry) ([]*stream, error) {
    38  	prefix := fmt.Sprintf("plugin.choria.adapter.%s.stream.", name)
    39  
    40  	instances, err := strconv.Atoi(cfg.Option(prefix+"workers", "10"))
    41  	if err != nil {
    42  		return nil, fmt.Errorf("%s should be a integer number", prefix+"workers")
    43  	}
    44  
    45  	servers := cfg.Option(prefix+"servers", "")
    46  
    47  	topic := cfg.Option(prefix+"topic", "")
    48  	if topic == "" {
    49  		topic = name
    50  	}
    51  
    52  	var workers []*stream
    53  
    54  	logger.Infof("Creating Choria Streams Adapter %s with %d workers publishing to %s", name, instances, topic)
    55  	for i := 0; i < instances; i++ {
    56  		logger.Debugf("Creating Choria Streams Adapter %s instance %d / %d publishing to message set %s", name, i, instances, topic)
    57  
    58  		iname := fmt.Sprintf("%s_%d-%s", name, i, strings.Replace(util.UniqueID(), "-", "", -1))
    59  
    60  		st := &stream{
    61  			clientID:    iname,
    62  			topic:       topic,
    63  			name:        fmt.Sprintf("%s.%d", name, i),
    64  			adapterName: name,
    65  			work:        work,
    66  			log:         logger.WithFields(logrus.Fields{"side": "stream", "instance": i}),
    67  		}
    68  
    69  		if servers != "" {
    70  			st.servers = st.resolver(strings.Split(servers, ","))
    71  		} else {
    72  			st.log.Warnf("%s not set, using standard client middleware resolution", prefix+"servers")
    73  			st.servers = fw.MiddlewareServers
    74  		}
    75  
    76  		workers = append(workers, st)
    77  	}
    78  
    79  	return workers, nil
    80  }
    81  
    82  func (sc *stream) resolver(parts []string) func() (srvcache.Servers, error) {
    83  	servers, err := srvcache.StringHostsToServers(parts, "nats")
    84  	return func() (srvcache.Servers, error) {
    85  		return servers, err
    86  	}
    87  }
    88  
    89  func (sc *stream) connect(ctx context.Context, cm inter.ConnectionManager) error {
    90  	if ctx.Err() != nil {
    91  		return fmt.Errorf("shutdown called")
    92  	}
    93  
    94  	nc, err := fw.NewConnector(ctx, sc.servers, sc.clientID, sc.log)
    95  	if err != nil {
    96  		return fmt.Errorf("could not start Choria Streams connection: %s", err)
    97  	}
    98  
    99  	sc.conn = nc
   100  
   101  	sc.log.Debugf("%s connected to Choria Streams", sc.clientID)
   102  
   103  	return nil
   104  }
   105  
   106  func (sc *stream) disconnect() {
   107  	if sc.conn != nil {
   108  		sc.log.Debugf("Disconnecting from Choria Streams")
   109  		sc.conn.Close()
   110  	}
   111  }
   112  
   113  func (sc *stream) publisher(ctx context.Context, wg *sync.WaitGroup) {
   114  	defer wg.Done()
   115  
   116  	bytes := stats.BytesCtr.WithLabelValues(sc.name, "output", cfg.Identity)
   117  	ectr := stats.ErrorCtr.WithLabelValues(sc.name, "output", cfg.Identity)
   118  	ctr := stats.ReceivedMsgsCtr.WithLabelValues(sc.name, "output", cfg.Identity)
   119  	timer := stats.ProcessTime.WithLabelValues(sc.name, "output", cfg.Identity)
   120  	workqlen := stats.WorkQueueLengthGauge.WithLabelValues(sc.adapterName, cfg.Identity)
   121  
   122  	transformerf := func(r ingest.Adaptable) {
   123  		obs := prometheus.NewTimer(timer)
   124  		defer obs.ObserveDuration()
   125  		defer func() { workqlen.Set(float64(len(sc.work))) }()
   126  
   127  		j, err := json.Marshal(transformer.TransformToOutput(r, "choria_streams"))
   128  		if err != nil {
   129  			sc.log.Warnf("Cannot JSON encode message for publishing to Choria Streams, discarding: %s", err)
   130  			ectr.Inc()
   131  			return
   132  		}
   133  
   134  		sc.log.Debugf("Publishing registration data from %s to %s", r.SenderID(), sc.topic)
   135  
   136  		bytes.Add(float64(len(j)))
   137  
   138  		err = sc.conn.PublishRaw(strings.ReplaceAll(sc.topic, "%s", r.SenderID()), j)
   139  		if err != nil {
   140  			sc.log.Warnf("Could not publish message to Choria Streams %s, discarding: %s", sc.topic, err)
   141  			ectr.Inc()
   142  			return
   143  		}
   144  
   145  		ctr.Inc()
   146  	}
   147  
   148  	for {
   149  		select {
   150  		case r := <-sc.work:
   151  			transformerf(r)
   152  
   153  		case <-ctx.Done():
   154  			sc.disconnect()
   155  
   156  			return
   157  		}
   158  	}
   159  }