github.com/hernad/nomad@v1.6.112/nomad/stream/ndjson.go (about)

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