github.com/aavshr/aws-sdk-go@v1.41.3/private/protocol/eventstream/eventstreamapi/stream_writer.go (about)

     1  package eventstreamapi
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/aavshr/aws-sdk-go/aws"
     9  )
    10  
    11  // StreamWriter provides concurrent safe writing to an event stream.
    12  type StreamWriter struct {
    13  	eventWriter *EventWriter
    14  	stream      chan eventWriteAsyncReport
    15  
    16  	done      chan struct{}
    17  	closeOnce sync.Once
    18  	err       *OnceError
    19  
    20  	streamCloser io.Closer
    21  }
    22  
    23  // NewStreamWriter returns a StreamWriter for the event writer, and stream
    24  // closer provided.
    25  func NewStreamWriter(eventWriter *EventWriter, streamCloser io.Closer) *StreamWriter {
    26  	w := &StreamWriter{
    27  		eventWriter:  eventWriter,
    28  		streamCloser: streamCloser,
    29  		stream:       make(chan eventWriteAsyncReport),
    30  		done:         make(chan struct{}),
    31  		err:          NewOnceError(),
    32  	}
    33  	go w.writeStream()
    34  
    35  	return w
    36  }
    37  
    38  // Close terminates the writers ability to write new events to the stream. Any
    39  // future call to Send will fail with an error.
    40  func (w *StreamWriter) Close() error {
    41  	w.closeOnce.Do(w.safeClose)
    42  	return w.Err()
    43  }
    44  
    45  func (w *StreamWriter) safeClose() {
    46  	close(w.done)
    47  }
    48  
    49  // ErrorSet returns a channel which will be closed
    50  // if an error occurs.
    51  func (w *StreamWriter) ErrorSet() <-chan struct{} {
    52  	return w.err.ErrorSet()
    53  }
    54  
    55  // Err returns any error that occurred while attempting to write an event to the
    56  // stream.
    57  func (w *StreamWriter) Err() error {
    58  	return w.err.Err()
    59  }
    60  
    61  // Send writes a single event to the stream returning an error if the write
    62  // failed.
    63  //
    64  // Send may be called concurrently. Events will be written to the stream
    65  // safely.
    66  func (w *StreamWriter) Send(ctx aws.Context, event Marshaler) error {
    67  	if err := w.Err(); err != nil {
    68  		return err
    69  	}
    70  
    71  	resultCh := make(chan error)
    72  	wrapped := eventWriteAsyncReport{
    73  		Event:  event,
    74  		Result: resultCh,
    75  	}
    76  
    77  	select {
    78  	case w.stream <- wrapped:
    79  	case <-ctx.Done():
    80  		return ctx.Err()
    81  	case <-w.done:
    82  		return fmt.Errorf("stream closed, unable to send event")
    83  	}
    84  
    85  	select {
    86  	case err := <-resultCh:
    87  		return err
    88  	case <-ctx.Done():
    89  		return ctx.Err()
    90  	case <-w.done:
    91  		return fmt.Errorf("stream closed, unable to send event")
    92  	}
    93  }
    94  
    95  func (w *StreamWriter) writeStream() {
    96  	defer w.Close()
    97  
    98  	for {
    99  		select {
   100  		case wrapper := <-w.stream:
   101  			err := w.eventWriter.WriteEvent(wrapper.Event)
   102  			wrapper.ReportResult(w.done, err)
   103  			if err != nil {
   104  				w.err.SetError(err)
   105  				return
   106  			}
   107  
   108  		case <-w.done:
   109  			if err := w.streamCloser.Close(); err != nil {
   110  				w.err.SetError(err)
   111  			}
   112  			return
   113  		}
   114  	}
   115  }
   116  
   117  type eventWriteAsyncReport struct {
   118  	Event  Marshaler
   119  	Result chan<- error
   120  }
   121  
   122  func (e eventWriteAsyncReport) ReportResult(cancel <-chan struct{}, err error) bool {
   123  	select {
   124  	case e.Result <- err:
   125  		return true
   126  	case <-cancel:
   127  		return false
   128  	}
   129  }