diff --git a/api/server/router/local.go b/api/server/router/local.go index 99db4242b3..7cb2a5a2f3 100644 --- a/api/server/router/local.go +++ b/api/server/router/local.go @@ -1,6 +1,11 @@ package router -import "github.com/docker/docker/api/server/httputils" +import ( + "net/http" + + "github.com/docker/docker/api/server/httputils" + "golang.org/x/net/context" +) // localRoute defines an individual API route to connect // with the docker daemon. It implements Route. @@ -59,3 +64,33 @@ func NewOptionsRoute(path string, handler httputils.APIFunc) Route { func NewHeadRoute(path string, handler httputils.APIFunc) Route { return NewRoute("HEAD", path, handler) } + +func cancellableHandler(h httputils.APIFunc) httputils.APIFunc { + return func(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error { + if notifier, ok := w.(http.CloseNotifier); ok { + notify := notifier.CloseNotify() + notifyCtx, cancel := context.WithCancel(ctx) + finished := make(chan struct{}) + defer close(finished) + ctx = notifyCtx + go func() { + select { + case <-notify: + cancel() + case <-finished: + } + }() + } + return h(ctx, w, r, vars) + } +} + +// Cancellable makes new route which embeds http.CloseNotifier feature to +// context.Context of handler. +func Cancellable(r Route) Route { + return localRoute{ + method: r.Method(), + path: r.Path(), + handler: cancellableHandler(r.Handler()), + } +}