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

     1  // Copyright (c) 2021-2022, 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  	"fmt"
    10  	"strconv"
    11  	"sync"
    12  
    13  	"github.com/choria-io/go-choria/broker/adapter/ingest"
    14  	"github.com/choria-io/go-choria/broker/adapter/stats"
    15  	"github.com/choria-io/go-choria/config"
    16  	"github.com/choria-io/go-choria/inter"
    17  	"github.com/sirupsen/logrus"
    18  )
    19  
    20  // Streams is an adapter that connects a NATS topic with messages sent from Choria
    21  // in its usual transport protocol to a Choria Streams stream.
    22  //
    23  // On the stream the messages will be JSON format with keys
    24  // body, sender and time.  Body is a base64 encoded string
    25  //
    26  // Configure the adapters:
    27  //
    28  //	# required
    29  //	plugin.choria.adapters = discovery
    30  //	plugin.choria.adapter.discovery.type = choria_streams
    31  //	plugin.choria.adapter.discovery.queue_len = 1000 # default
    32  //
    33  // Configure the stream output:
    34  //
    35  //	plugin.choria.adapter.discovery.stream.servers = js1:4222,js2:4222 # uses normal middleware server resolution when unset
    36  //	plugin.choria.adapter.discovery.stream.topic = discovery # default, %s gets replaced with sender id
    37  //	plugin.choria.adapter.discovery.stream.workers = 10 # default
    38  //
    39  // Configure the NATS ingest:
    40  //
    41  //	plugin.choria.adapter.discovery.ingest.topic = mcollective.broadcast.agent.discovery
    42  //	plugin.choria.adapter.discovery.ingest.protocol = request # or reply
    43  //	plugin.choria.adapter.discovery.ingest.workers = 10 # default
    44  type Streams struct {
    45  	streams []*stream
    46  	ingests []*ingest.NatsIngest
    47  	work    chan ingest.Adaptable
    48  	log     *logrus.Entry
    49  }
    50  
    51  var fw inter.Framework
    52  var cfg *config.Config
    53  
    54  func Create(name string, choria inter.Framework) (adapter *Streams, err error) {
    55  	fw = choria
    56  	cfg = fw.Configuration()
    57  
    58  	s := fmt.Sprintf("plugin.choria.adapter.%s.queue_len", name)
    59  	worklen, err := strconv.Atoi(cfg.Option(s, "1000"))
    60  	if err != nil {
    61  		return nil, fmt.Errorf("%s should be a integer number", s)
    62  	}
    63  
    64  	stats.WorkQueueCapacityGauge.WithLabelValues(name, cfg.Identity).Set(float64(worklen))
    65  
    66  	adapter = &Streams{
    67  		log:  fw.Logger("streams_adapter").WithFields(logrus.Fields{"name": name}),
    68  		work: make(chan ingest.Adaptable, worklen),
    69  	}
    70  
    71  	adapter.ingests, err = ingest.New(name, adapter.work, choria, adapter.log)
    72  	if err != nil {
    73  		return nil, fmt.Errorf("could not create adapter %s: %s", name, err)
    74  	}
    75  
    76  	adapter.streams, err = newStream(name, adapter.work, adapter.log)
    77  	if err != nil {
    78  		return nil, fmt.Errorf("could not create adapter %s: %s", name, err)
    79  	}
    80  
    81  	return adapter, nil
    82  }
    83  
    84  func (sa *Streams) Init(ctx context.Context, cm inter.ConnectionManager) (err error) {
    85  	for _, worker := range sa.streams {
    86  		if ctx.Err() != nil {
    87  			return fmt.Errorf("shutdown called")
    88  		}
    89  
    90  		err = worker.connect(ctx, cm)
    91  		if err != nil {
    92  			return fmt.Errorf("failure during initial Choria Streams connections: %s", err)
    93  		}
    94  	}
    95  
    96  	for _, worker := range sa.ingests {
    97  		if ctx.Err() != nil {
    98  			return fmt.Errorf("shutdown called")
    99  		}
   100  
   101  		err = worker.Connect(ctx, cm)
   102  		if err != nil {
   103  			return fmt.Errorf("failure during Choria Streams initial connections: %s", err)
   104  		}
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (sa *Streams) Process(ctx context.Context, wg *sync.WaitGroup) {
   111  	defer wg.Done()
   112  
   113  	for _, worker := range sa.streams {
   114  		wg.Add(1)
   115  		go worker.publisher(ctx, wg)
   116  	}
   117  
   118  	for _, worker := range sa.ingests {
   119  		wg.Add(1)
   120  		go worker.Receiver(ctx, wg)
   121  	}
   122  }