func TestPublisherReconnect(t *testing.T) {
server := httptest.NewServer(baseServer.router())
defer server.Close()
client := &http.Client{Transport: &http.Transport{}}
uuid, _ := util.NewUUID()
url := server.URL + "/streams/" + uuid
// curl -XPUT <url>/streams/<uuid>
request, _ := http.NewRequest("PUT", url, nil)
resp, err := client.Do(request)
defer resp.Body.Close()
assert.Nil(t, err)
done := make(chan bool)
writer, err := broker.NewWriter(uuid)
assert.Nil(t, err)
_, err = writer.Write([]byte("hello"))
assert.Nil(t, err)
defer writer.Close()
go func() {
// curl <url>/streams/<uuid>
// -- waiting for publish to arrive
resp, err := http.Get(server.URL + "/streams/" + uuid)
defer resp.Body.Close()
assert.Nil(t, err)
body, _ := ioutil.ReadAll(resp.Body)
assert.Equal(t, body, []byte("hello world"))
done <- true
}()
// curl -XPOST -H "Transfer-Encoding: chunked" -d "hello" <url>/streams/<uuid>
req, _ := http.NewRequest("POST", server.URL+"/streams/"+uuid, bytes.NewReader([]byte("hello world")))
req.TransferEncoding = []string{"chunked"}
r, err := client.Do(req)
r.Body.Close()
assert.Nil(t, err)
<-done
// Read the whole response after the publisher has
// completed. The mechanics of this is different in that
// most of the content will be replayed instead of received
// in chunks as they arrive.
resp, err = http.Get(server.URL + "/streams/" + uuid)
defer resp.Body.Close()
assert.Nil(t, err)
body, _ := ioutil.ReadAll(resp.Body)
assert.Equal(t, body, []byte("hello world"))
}
func (s *Server) publish(w http.ResponseWriter, r *http.Request) {
if !util.StringInSlice(r.TransferEncoding, "chunked") {
http.Error(w, "A chunked Transfer-Encoding header is required.", http.StatusBadRequest)
return
}
writer, err := broker.NewWriter(key(r))
if err != nil {
handleError(w, r, err)
return
}
body := bufio.NewReader(r.Body)
defer r.Body.Close()
wl, err := broker.Len(writer)
if err != nil {
handleError(w, r, err)
return
}
if wl > 0 {
_, err = body.Discard(int(wl))
if err != nil {
handleError(w, r, err)
return
}
}
_, err = io.Copy(writer, body)
if err == io.ErrUnexpectedEOF {
util.CountWithData("server.pub.read.eoferror", 1, "msg=%q request_id=%q", err, r.Header.Get("Request-Id"))
return
}
netErr, ok := err.(net.Error)
if ok && netErr.Timeout() {
util.CountWithData("server.pub.read.timeout", 1, "msg=%q request_id=%q", err, r.Header.Get("Request-Id"))
handleError(w, r, netErr)
return
}
if err != nil {
log.Printf("%#v", err)
http.Error(w, "Unhandled error, please try again.", http.StatusInternalServerError)
rollbar.Error(rollbar.ERR, fmt.Errorf("unhandled error: %#v", err))
return
}
util.CountWithData("server.pub.read.end", 1, "request_id=%q", r.Header.Get("Request-Id"))
writer.Close()
// Asynchronously upload the output to our defined storage backend.
go storeOutput(key(r), requestURI(r), s.StorageBaseURL)
}
func pub(w http.ResponseWriter, r *http.Request) {
if !util.StringSliceUtil(r.TransferEncoding).Contains("chunked") {
http.Error(w, "A chunked Transfer-Encoding header is required.", http.StatusBadRequest)
return
}
writer, err := broker.NewWriter(key(r))
if err != nil {
handleError(w, r, err)
return
}
defer writer.Close()
body := bufio.NewReader(r.Body)
defer r.Body.Close()
done := make(chan struct{})
go func() {
for {
select {
case <-done:
// Asynchronously upload the output to our defined storage backend.
go storeOutput(key(r), requestURI(r))
return
case <-time.After(*util.StorageInterval):
// Asynchronously upload the output to our defined storage backend.
go storeOutput(key(r), requestURI(r))
}
}
}()
_, err = io.Copy(writer, body)
if err == io.ErrUnexpectedEOF {
util.CountWithData("server.pub.read.eoferror", 1, "msg=\"%v\"", err.Error())
return
}
if err != nil {
log.Printf("%#v", err)
http.Error(w, "Unhandled error, please try again.", http.StatusInternalServerError)
rollbar.Error(rollbar.ERR, fmt.Errorf("unhandled error: %#v", err))
return
}
close(done)
}
func TestPubClosed(t *testing.T) {
uuid, _ := util.NewUUID()
registrar := broker.NewRedisRegistrar()
err := registrar.Register(uuid)
assert.Nil(t, err)
writer, err := broker.NewWriter(uuid)
assert.Nil(t, err)
writer.Close()
request, _ := http.NewRequest("POST", "/streams/"+uuid, nil)
request.TransferEncoding = []string{"chunked"}
response := httptest.NewRecorder()
baseServer.publish(response, request)
assert.Equal(t, response.Code, http.StatusNotFound)
}