github.com/kubeshop/testkube@v1.17.23/pkg/logs/client/stream.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/nats-io/nats.go"
    10  	"github.com/nats-io/nats.go/jetstream"
    11  	"go.uber.org/zap"
    12  
    13  	"github.com/kubeshop/testkube/pkg/log"
    14  	"github.com/kubeshop/testkube/pkg/logs/events"
    15  	"github.com/kubeshop/testkube/pkg/utils"
    16  )
    17  
    18  const ConsumerPrefix = "lc"
    19  
    20  func NewNatsLogStream(nc *nats.Conn) (s Stream, err error) {
    21  	js, err := jetstream.New(nc)
    22  	if err != nil {
    23  		return s, err
    24  	}
    25  
    26  	return &NatsLogStream{
    27  		nc:  nc,
    28  		js:  js,
    29  		log: log.DefaultLogger,
    30  	}, nil
    31  }
    32  
    33  type NatsLogStream struct {
    34  	nc  *nats.Conn
    35  	js  jetstream.JetStream
    36  	log *zap.SugaredLogger
    37  }
    38  
    39  func (c NatsLogStream) Init(ctx context.Context, id string) (StreamMetadata, error) {
    40  	s, err := c.js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
    41  		Name:    c.Name(id),
    42  		Storage: jetstream.FileStorage, // durable stream
    43  	})
    44  
    45  	if err == nil {
    46  		c.log.Debugw("stream upserted", "info", s.CachedInfo())
    47  	}
    48  
    49  	return StreamMetadata{Name: c.Name(id)}, err
    50  
    51  }
    52  
    53  func (c NatsLogStream) Finish(ctx context.Context, id string) error {
    54  	return c.Push(ctx, id, events.NewFinishLog())
    55  }
    56  
    57  // Push log chunk to NATS stream
    58  func (c NatsLogStream) Push(ctx context.Context, id string, log *events.Log) error {
    59  	b, err := json.Marshal(log)
    60  	if err != nil {
    61  		return err
    62  	}
    63  	return c.PushBytes(ctx, id, b)
    64  }
    65  
    66  // Push log chunk to NATS stream
    67  // TODO handle message repeat with backoff strategy on error
    68  func (c NatsLogStream) PushBytes(ctx context.Context, id string, bytes []byte) error {
    69  	_, err := c.js.Publish(ctx, c.Name(id), bytes)
    70  	return err
    71  }
    72  
    73  // Start emits start event to the stream - logs service will handle start and create new stream
    74  func (c NatsLogStream) Start(ctx context.Context, id string) (resp StreamResponse, err error) {
    75  	return c.syncCall(ctx, StartSubject, id)
    76  }
    77  
    78  // Stop emits stop event to the stream and waits for given stream to be stopped fully - logs service will handle stop and close stream and all subscribers
    79  func (c NatsLogStream) Stop(ctx context.Context, id string) (resp StreamResponse, err error) {
    80  	return c.syncCall(ctx, StopSubject, id)
    81  }
    82  
    83  // Get returns channel with log stream chunks for given execution id connects through GRPC to log service
    84  func (c NatsLogStream) Get(ctx context.Context, id string) (chan events.LogResponse, error) {
    85  	ch := make(chan events.LogResponse)
    86  
    87  	name := fmt.Sprintf("%s%s%s", ConsumerPrefix, id, utils.RandAlphanum(6))
    88  	cons, err := c.js.CreateOrUpdateConsumer(
    89  		ctx,
    90  		c.Name(id),
    91  		jetstream.ConsumerConfig{
    92  			Name:          name,
    93  			Durable:       name,
    94  			DeliverPolicy: jetstream.DeliverAllPolicy,
    95  		},
    96  	)
    97  
    98  	if err != nil {
    99  		return ch, err
   100  	}
   101  
   102  	log := c.log.With("id", id)
   103  
   104  	go func() {
   105  		defer close(ch)
   106  		for {
   107  			msg, err := cons.Next()
   108  			if err != nil {
   109  				ch <- events.LogResponse{Error: err}
   110  				return
   111  			}
   112  
   113  			if finished := c.handleJetstreamMessage(log, ch, msg); finished {
   114  				return
   115  			}
   116  		}
   117  	}()
   118  
   119  	return ch, nil
   120  }
   121  
   122  func (c NatsLogStream) handleJetstreamMessage(log *zap.SugaredLogger, ch chan events.LogResponse, msg jetstream.Msg) (finish bool) {
   123  	// deliver to subscriber
   124  	logChunk := events.Log{}
   125  	err := json.Unmarshal(msg.Data(), &logChunk)
   126  	if err != nil {
   127  		if err := msg.Nak(); err != nil {
   128  			log.Errorw("error nacking message", "error", err)
   129  			ch <- events.LogResponse{Error: err}
   130  			return
   131  		}
   132  		return
   133  	}
   134  
   135  	if err := msg.Ack(); err != nil {
   136  		ch <- events.LogResponse{Error: err}
   137  		log.Errorw("error acking message", "error", err)
   138  		return
   139  	}
   140  
   141  	if events.IsFinished(&logChunk) {
   142  		return true
   143  	}
   144  
   145  	ch <- events.LogResponse{Log: logChunk}
   146  
   147  	return
   148  }
   149  
   150  // syncCall sends request to given subject and waits for response
   151  func (c NatsLogStream) syncCall(ctx context.Context, subject, id string) (resp StreamResponse, err error) {
   152  	b, err := json.Marshal(events.NewTrigger(id))
   153  	if err != nil {
   154  		return resp, err
   155  	}
   156  	m, err := c.nc.Request(subject, b, time.Minute)
   157  	if err != nil {
   158  		return resp, err
   159  	}
   160  
   161  	return StreamResponse{Message: m.Data}, nil
   162  }
   163  
   164  func (c NatsLogStream) Name(id ...string) string {
   165  	if len(id) > 0 {
   166  		return StreamPrefix + id[0]
   167  	}
   168  
   169  	return StreamPrefix + utils.RandAlphanum(10)
   170  }