github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/nomad/stream/ndjson.go (about)

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