This commit is contained in:
Alex Dadgar 2016-07-18 18:41:21 -07:00
parent b2df901158
commit 0327109300
2 changed files with 121 additions and 9 deletions

View file

@ -261,7 +261,6 @@ func (s *StreamFramer) Destroy() {
s.l.Lock()
wasRunning := s.running
s.running = false
s.f = nil
close(s.shutdownCh)
s.heartbeat.Stop()
s.flusher.Stop()
@ -271,6 +270,7 @@ func (s *StreamFramer) Destroy() {
if wasRunning {
<-s.exitCh
}
s.out.Close()
}
// Run starts a long lived goroutine that handles sending data as well as
@ -299,7 +299,6 @@ func (s *StreamFramer) run() {
defer func() {
s.l.Lock()
s.err = err
s.out.Close()
close(s.exitCh)
close(s.outbound)
s.l.Unlock()
@ -322,8 +321,11 @@ func (s *StreamFramer) run() {
// Read the data for the frame, and send it
s.f.Data = s.readData()
s.outbound <- s.f
select {
case s.outbound <- s.f:
s.f = nil
default:
}
s.l.Unlock()
case <-s.heartbeat.C:
@ -339,7 +341,7 @@ OUTER:
case <-s.shutdownCh:
break OUTER
case o := <-s.outbound:
// Send the frame and then clear the current working frame
// Send the frame
if err = s.enc.Encode(o); err != nil {
return
}

View file

@ -1,6 +1,7 @@
package agent
import (
"bytes"
"io"
"io/ioutil"
"net/http"
@ -8,6 +9,7 @@ import (
"os"
"path/filepath"
"reflect"
"strconv"
"testing"
"time"
@ -302,6 +304,98 @@ func TestStreamFramer_Heartbeat(t *testing.T) {
}
}
// This test checks that frames are received in order
func TestStreamFramer_Order(t *testing.T) {
// Create the stream framer
r, w := io.Pipe()
wrappedW := &WriteCloseChecker{WriteCloser: w}
// Ensure the batch window doesn't get hit
hRate, bWindow := 100*time.Millisecond, 100*time.Millisecond
sf := NewStreamFramer(wrappedW, hRate, bWindow, 100)
sf.Run()
// Create a decoder
dec := codec.NewDecoder(r, jsonHandle)
//files := []string{"1", "2", "3", "4", "5"}
files := []string{"1"}
input := bytes.NewBuffer(make([]byte, 100000))
for i := 0; i <= 2000; i++ {
str := strconv.Itoa(i) + "\n"
input.WriteString(str)
}
expected := bytes.NewBuffer(make([]byte, 100000))
for _, _ = range files {
expected.Write(input.Bytes())
}
receivedBuf := bytes.NewBuffer(make([]byte, 100000))
// Start the reader
resultCh := make(chan struct{})
go func() {
for {
var frame StreamFrame
if err := dec.Decode(&frame); err != nil {
t.Fatalf("failed to decode")
}
if frame.IsHeartbeat() {
continue
}
receivedBuf.Write(frame.Data)
if reflect.DeepEqual(expected, receivedBuf) {
resultCh <- struct{}{}
return
}
}
}()
// Send the data
b := input.Bytes()
shards := 10
each := len(b) / shards
for _, f := range files {
for i := 0; i < shards; i++ {
l, r := each*i, each*(i+1)
if i == shards-1 {
r = len(b)
}
if err := sf.Send(f, "", b[l:r], 0); err != nil {
t.Fatalf("Send() failed %v", err)
}
}
}
// Ensure we get data
select {
case <-resultCh:
case <-time.After(4 * bWindow):
got := receivedBuf.String()
want := expected.String()
t.Fatalf("Did not receive data in sorted order\nGot:%v\nWant:%v\n", got, want)
}
// Close the reader and wait. This should cause the runner to exit
if err := r.Close(); err != nil {
t.Fatalf("failed to close reader")
}
select {
case <-sf.ExitCh():
case <-time.After(2 * hRate):
t.Fatalf("exit channel should close")
}
sf.Destroy()
if !wrappedW.Closed {
t.Fatalf("writer not closed")
}
}
func TestHTTP_Stream_MissingParams(t *testing.T) {
httpTest(t, nil, func(s *TestServer) {
req, err := http.NewRequest("GET", "/v1/client/fs/stream/", nil)
@ -364,7 +458,11 @@ func TestHTTP_Stream_NoFile(t *testing.T) {
ad := tempAllocDir(t)
defer os.RemoveAll(ad.AllocDir)
if err := s.Server.stream(0, "foo", ad, nopWriteCloser{ioutil.Discard}); err == nil {
framer := NewStreamFramer(nopWriteCloser{ioutil.Discard}, streamHeartbeatRate, streamBatchWindow, streamFrameSize)
framer.Run()
defer framer.Destroy()
if err := s.Server.stream(0, "foo", ad, framer, nil); err == nil {
t.Fatalf("expected an error when streaming unknown file")
}
})
@ -419,9 +517,13 @@ func TestHTTP_Stream_Modify(t *testing.T) {
t.Fatalf("write failed: %v", err)
}
framer := NewStreamFramer(w, streamHeartbeatRate, streamBatchWindow, streamFrameSize)
framer.Run()
defer framer.Destroy()
// Start streaming
go func() {
if err := s.Server.stream(0, streamFile, ad, w); err != nil {
if err := s.Server.stream(0, streamFile, ad, framer, nil); err != nil {
t.Fatalf("stream() failed: %v", err)
}
}()
@ -496,9 +598,13 @@ func TestHTTP_Stream_Truncate(t *testing.T) {
t.Fatalf("write failed: %v", err)
}
framer := NewStreamFramer(w, streamHeartbeatRate, streamBatchWindow, streamFrameSize)
framer.Run()
defer framer.Destroy()
// Start streaming
go func() {
if err := s.Server.stream(0, streamFile, ad, w); err != nil {
if err := s.Server.stream(0, streamFile, ad, framer, nil); err != nil {
t.Fatalf("stream() failed: %v", err)
}
}()
@ -595,9 +701,12 @@ func TestHTTP_Stream_Delete(t *testing.T) {
t.Fatalf("write failed: %v", err)
}
framer := NewStreamFramer(wrappedW, streamHeartbeatRate, streamBatchWindow, streamFrameSize)
framer.Run()
// Start streaming
go func() {
if err := s.Server.stream(0, streamFile, ad, wrappedW); err != nil {
if err := s.Server.stream(0, streamFile, ad, framer, nil); err != nil {
t.Fatalf("stream() failed: %v", err)
}
}()
@ -615,6 +724,7 @@ func TestHTTP_Stream_Delete(t *testing.T) {
t.Fatalf("did not receive delete")
}
framer.Destroy()
testutil.WaitForResult(func() (bool, error) {
return wrappedW.Closed, nil
}, func(err error) {