frames to reader

api/fs.go

@@ -204,7 +204,6 @@ func (a *AllocFS) getErrorMsg(resp *http.Response) error {
 // * path: path to file to stream.
 // * offset: The offset to start streaming data at.
 // * origin: Either "start" or "end" and defines from where the offset is applied.
-// * cancel: A channel which when closed will stop streaming.
 //
 // The return value is a channel that will emit StreamFrames as they are read.
 func (a *AllocFS) Stream(alloc *Allocation, path, origin string, offset int64,
@@ -275,3 +274,82 @@ func (a *AllocFS) Stream(alloc *Allocation, path, origin string, offset int64,
 
 	return frames, nil, nil
 }
+
+// FrameReader is used to convert a stream of frames into a read closer.
+type FrameReader struct {
+	frames   <-chan *StreamFrame
+	cancelCh chan struct{}
+
+	frame       *StreamFrame
+	frameOffset int
+
+	// To handle printing the file events
+	fileEventOffset int
+	fileEvent       []byte
+
+	byteOffset int
+}
+
+// NewFrameReader takes a channel of frames and returns a FrameReader which
+// implements io.ReadCloser
+func NewFrameReader(frames <-chan *StreamFrame, cancelCh chan struct{}) *FrameReader {
+	return &FrameReader{
+		frames:   frames,
+		cancelCh: cancelCh,
+	}
+}
+
+// Offset returns the offset into the stream.
+func (f *FrameReader) Offset() int {
+	return f.byteOffset
+}
+
+// Read reads the data of the incoming frames into the bytes buffer. Returns EOF
+// when there are no more frames.
+func (f *FrameReader) Read(p []byte) (n int, err error) {
+	if f.frame == nil {
+		frame, ok := <-f.frames
+		if !ok {
+			return 0, io.EOF
+		}
+		f.frame = frame
+	}
+
+	if f.frame.FileEvent != "" && len(f.fileEvent) == 0 {
+		f.fileEvent = []byte(fmt.Sprintf("\nnomad: %q\n", f.frame.FileEvent))
+		f.fileEventOffset = 0
+	}
+
+	// If there is a file event we inject it into the read stream
+	if l := len(f.fileEvent); l != 0 && l != f.fileEventOffset {
+		n = copy(p, f.fileEvent[f.fileEventOffset:])
+		f.fileEventOffset += n
+		return n, nil
+	}
+
+	if len(f.fileEvent) == f.fileEventOffset {
+		f.fileEvent = nil
+		f.fileEventOffset = 0
+	}
+
+	// Copy the data out of the frame and update our offset
+	n = copy(p, f.frame.Data[f.frameOffset:])
+	f.frameOffset += n
+
+	// Store the total offset into the file
+	f.byteOffset = int(f.frame.Offset) + f.frameOffset
+
+	// Clear the frame and its offset once we have read everything
+	if len(f.frame.Data) == f.frameOffset {
+		f.frame = nil
+		f.frameOffset = 0
+	}
+
+	return n, nil
+}
+
+// Close cancels the stream of frames
+func (f *FrameReader) Close() error {
+	close(f.cancelCh)
+	return nil
+}

command/fs.go

@@ -312,43 +312,20 @@ func (f *FSCommand) followFile(client *api.Client, alloc *api.Allocation,
 	signalCh := make(chan os.Signal, 1)
 	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
 
-	var frame *api.StreamFrame
-	var ok bool
-	for {
-		select {
-		case <-signalCh:
-			// End the streaming
-			close(cancel)
+	// Create a reader
+	r := api.NewFrameReader(frames, cancel)
 
-			// Output the last offset
-			if frame != nil && frame.Offset > 0 {
-				f.Ui.Output(fmt.Sprintf("\nLast outputted offset (bytes): %d", frame.Offset))
-			}
+	go func() {
+		<-signalCh
 
-			return nil
-		case frame, ok = <-frames:
-			if !ok {
-				// Connection has been killed
-				return nil
-			}
+		// End the streaming
+		r.Close()
 
-			if frame == nil {
-				panic("received nil frame; please report as a bug")
-			}
-
-			if frame.IsHeartbeat() {
-				continue
-			}
-
-			// Print the file event
-			if frame.FileEvent != "" {
-				f.Ui.Output(fmt.Sprintf("nomad: FileEvent %q", frame.FileEvent))
-			}
-
-			fmt.Print(string(frame.Data))
-		}
-	}
+		// Output the last offset
+		f.Ui.Output(fmt.Sprintf("\nLast outputted offset (bytes): %d", r.Offset()))
+	}()
 
+	io.Copy(os.Stdout, r)
 	return nil
 }