github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/nomad/stream/ndjson.go (about)

     1  package stream
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/hashicorp/go-msgpack/codec"
    10  
    11  	"github.com/hashicorp/nomad/nomad/structs"
    12  )
    13  
    14  var (
    15  	// JsonHeartbeat is an empty JSON object to send as a heartbeat
    16  	// Avoids creating many heartbeat instances
    17  	JsonHeartbeat = &structs.EventJson{Data: []byte("{}")}
    18  )
    19  
    20  // JsonStream is used to send new line delimited JSON and heartbeats
    21  // to a destination (out channel)
    22  type JsonStream struct {
    23  	// ctx is a passed in context used to notify the json stream
    24  	// when it should terminate
    25  	ctx context.Context
    26  
    27  	outCh chan *structs.EventJson
    28  
    29  	// heartbeat is the interval to send heartbeat messages to keep a connection
    30  	// open.
    31  	heartbeatTick *time.Ticker
    32  }
    33  
    34  // NewJsonStream creates a new json stream that will output Json structs
    35  // to the passed output channel. The constructor starts a goroutine
    36  // to begin heartbeating on its set interval and also sends an initial heartbeat
    37  // to notify the client about the successful connection initialization.
    38  func NewJsonStream(ctx context.Context, heartbeat time.Duration) *JsonStream {
    39  	s := &JsonStream{
    40  		ctx:           ctx,
    41  		outCh:         make(chan *structs.EventJson, 10),
    42  		heartbeatTick: time.NewTicker(heartbeat),
    43  	}
    44  
    45  	s.outCh <- JsonHeartbeat
    46  	go s.heartbeat()
    47  
    48  	return s
    49  }
    50  
    51  func (n *JsonStream) OutCh() chan *structs.EventJson {
    52  	return n.outCh
    53  }
    54  
    55  func (n *JsonStream) heartbeat() {
    56  	for {
    57  		select {
    58  		case <-n.ctx.Done():
    59  			return
    60  		case <-n.heartbeatTick.C:
    61  			// Send a heartbeat frame
    62  			select {
    63  			case n.outCh <- JsonHeartbeat:
    64  			case <-n.ctx.Done():
    65  				return
    66  			}
    67  		}
    68  	}
    69  }
    70  
    71  // Send encodes an object into Newline delimited json. An error is returned
    72  // if json encoding fails or if the stream is no longer running.
    73  func (n *JsonStream) Send(v interface{}) error {
    74  	if n.ctx.Err() != nil {
    75  		return n.ctx.Err()
    76  	}
    77  
    78  	var buf bytes.Buffer
    79  	enc := codec.NewEncoder(&buf, structs.JsonHandleWithExtensions)
    80  	err := enc.Encode(v)
    81  	if err != nil {
    82  		return fmt.Errorf("error marshaling json for stream: %w", err)
    83  	}
    84  
    85  	select {
    86  	case <-n.ctx.Done():
    87  		return fmt.Errorf("error stream is no longer running: %w", err)
    88  	case n.outCh <- &structs.EventJson{Data: buf.Bytes()}:
    89  	}
    90  
    91  	return nil
    92  }