harness-drone/pkg/server/ws.go
2015-06-22 20:45:08 -07:00

135 lines
2.8 KiB
Go

package server
import (
"fmt"
"io"
"net/http"
"strconv"
"github.com/drone/drone/pkg/bus"
log "github.com/drone/drone/Godeps/_workspace/src/github.com/Sirupsen/logrus"
"github.com/drone/drone/Godeps/_workspace/src/github.com/gin-gonic/gin"
"github.com/drone/drone/Godeps/_workspace/src/github.com/manucorporat/sse"
"github.com/drone/drone/pkg/docker"
)
// GetRepoEvents will upgrade the connection to a Websocket and will stream
// event updates to the browser.
func GetRepoEvents(c *gin.Context) {
bus_ := ToBus(c)
repo := ToRepo(c)
c.Writer.Header().Set("Content-Type", "text/event-stream")
eventc := make(chan *bus.Event, 1)
bus_.Subscribe(eventc)
defer func() {
bus_.Unsubscribe(eventc)
close(eventc)
log.Infof("closed event stream")
}()
c.Stream(func(w io.Writer) bool {
select {
case event := <-eventc:
if event == nil {
log.Infof("nil event received")
return false
}
if event.Kind == bus.EventRepo &&
event.Name == repo.FullName {
sse.Encode(w, sse.Event{
Event: "message",
Data: string(event.Msg),
})
}
case <-c.Writer.CloseNotify():
return false
}
return true
})
}
func GetStream(c *gin.Context) {
conf := ToSettings(c)
store := ToDatastore(c)
repo := ToRepo(c)
runner := ToRunner(c)
commitseq, _ := strconv.Atoi(c.Params.ByName("build"))
jobnum, _ := strconv.Atoi(c.Params.ByName("number"))
c.Writer.Header().Set("Content-Type", "text/event-stream")
build, err := store.BuildNumber(repo, commitseq)
if err != nil {
c.Fail(404, err)
return
}
job, err := store.JobNumber(build, jobnum)
if err != nil {
c.Fail(404, err)
return
}
var rc io.ReadCloser
// if the commit is being executed by an agent
// we'll proxy the build output directly to the
// remote Docker client, through the agent.
if conf.Agents.Secret != "" {
addr, err := store.Agent(build)
if err != nil {
c.Fail(500, err)
return
}
url := fmt.Sprintf("http://%s/stream/%d?token=%s", addr, job.ID, conf.Agents.Secret)
resp, err := http.Get(url)
if err != nil {
c.Fail(500, err)
return
} else if resp.StatusCode != 200 {
resp.Body.Close()
c.AbortWithStatus(resp.StatusCode)
return
}
rc = resp.Body
} else {
// else if the commit is not being executed
// by the build agent we can use the local runner
rc, err = runner.Logs(job)
if err != nil {
c.Fail(404, err)
return
}
}
defer func() {
rc.Close()
}()
go func() {
<-c.Writer.CloseNotify()
rc.Close()
}()
rw := &StreamWriter{c.Writer, 0}
docker.StdCopy(rw, rw, rc)
}
type StreamWriter struct {
writer gin.ResponseWriter
count int
}
func (w *StreamWriter) Write(data []byte) (int, error) {
var err = sse.Encode(w.writer, sse.Event{
Id: strconv.Itoa(w.count),
Event: "message",
Data: string(data),
})
w.writer.Flush()
w.count += len(data)
return len(data), err
}