Handle cancellation

This commit is contained in:
criyle 2020-06-04 00:11:48 -04:00
parent 891ed8165b
commit f61a319cd6
7 changed files with 43 additions and 17 deletions

View File

@ -25,7 +25,7 @@ func handleRun(c *gin.Context) {
c.AbortWithStatusJSON(http.StatusBadRequest, err.Error)
return
}
rt := <-work.Submit(r)
rt := <-work.Submit(c.Request.Context(), r)
execObserve(rt)
if rt.Error != nil {
c.Error(rt.Error)

View File

@ -34,7 +34,7 @@ func (e *execServer) Exec(ctx context.Context, req *pb.Request) (*pb.Response, e
if len(si) > 0 || len(so) > 0 {
return nil, fmt.Errorf("Stream in / out are not avaliable for exec request")
}
rt := <-work.Submit(r)
rt := <-work.Submit(ctx, r)
execObserve(rt)
if rt.Error != nil {
return nil, err
@ -90,7 +90,7 @@ func (e *execServer) ExecStream(es pb.Executor_ExecStreamServer) error {
}
}
rtCh := work.Submit(rq)
rtCh := work.Execute(es.Context(), rq)
for {
select {
case err := <-errCh:

View File

@ -1,6 +1,7 @@
package main
import (
"context"
"net/http"
"time"
@ -37,6 +38,10 @@ func handleWS(c *gin.Context) {
conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
for {
req := new(model.Request)
if err := conn.ReadJSON(req); err != nil {
@ -49,7 +54,7 @@ func handleWS(c *gin.Context) {
return
}
go func() {
ret := <-work.Submit(r)
ret := <-work.Submit(ctx, r)
execObserve(ret)
resultCh <- model.ConvertResponse(ret)
}()

View File

@ -1,6 +1,7 @@
package envexec
import (
"context"
"fmt"
"golang.org/x/sync/errgroup"
@ -32,7 +33,7 @@ type Pipe struct {
}
// Run starts the cmd and returns exec results
func (r *Group) Run() ([]Result, error) {
func (r *Group) Run(ctx context.Context) ([]Result, error) {
// prepare files
fds, pipeToCollect, fileToClose, err := prepareFds(r)
defer func() { closeFiles(fileToClose) }()
@ -58,7 +59,7 @@ func (r *Group) Run() ([]Result, error) {
for i, c := range r.Cmd {
i, c := i, c
g.Go(func() error {
r, err := runSingle(ms[i], c, fds[i], pipeToCollect[i])
r, err := runSingle(ctx, ms[i], c, fds[i], pipeToCollect[i])
result[i] = r
if err != nil {
result[i].Status = StatusInternalError

View File

@ -8,7 +8,7 @@ import (
)
// runSingle runs Cmd inside the given environment and cgroup
func runSingle(m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (result Result, err error) {
func runSingle(pc context.Context, m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (result Result, err error) {
fdToClose := fds
defer func() { closeFiles(fdToClose) }()
@ -32,7 +32,7 @@ func runSingle(m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (resu
}
// start the cmd (they will be canceled in other goroutines)
ctx, cancel := context.WithCancel(context.TODO())
ctx, cancel := context.WithCancel(pc)
waiterCtx, waiterCancel := context.WithCancel(ctx)
process, err := m.Execve(ctx, execParam)

View File

@ -1,6 +1,7 @@
package envexec
import (
"context"
"fmt"
)
@ -15,7 +16,7 @@ type Single struct {
}
// Run starts the cmd and returns exec results
func (s *Single) Run() (result Result, err error) {
func (s *Single) Run(ctx context.Context) (result Result, err error) {
// prepare files
fd, fileToClose, pipeToCollect, err := prepareCmdFd(s.Cmd, len(s.Cmd.Files))
defer func() { closeFiles(fileToClose) }()
@ -31,7 +32,7 @@ func (s *Single) Run() (result Result, err error) {
}
defer s.EnvironmentPool.Put(m)
result, err = runSingle(m, s.Cmd, fd, pipeToCollect)
result, err = runSingle(ctx, m, s.Cmd, fd, pipeToCollect)
fileToClose = nil // already closed by runOne
if err != nil {
result.Status = StatusInternalError

View File

@ -1,6 +1,7 @@
package worker
import (
"context"
"fmt"
"path"
"sync"
@ -30,6 +31,7 @@ type Worker struct {
type workRequest struct {
*Request
context.Context
resultCh chan<- Response
}
@ -56,15 +58,32 @@ func (w *Worker) Start() {
}
// Submit submits a single request
func (w *Worker) Submit(req *Request) <-chan Response {
func (w *Worker) Submit(ctx context.Context, req *Request) <-chan Response {
ch := make(chan Response, 1)
w.workCh <- workRequest{
Request: req,
Context: ctx,
resultCh: ch,
}
return ch
}
// Execute will execute the request in new goroutine (bypass the parallism limit)
func (w *Worker) Execute(ctx context.Context, req *Request) <-chan Response {
ch := make(chan Response, 1)
w.wg.Add(1)
go func() {
defer w.wg.Done()
wq := workRequest{
Request: req,
Context: ctx,
resultCh: ch,
}
w.workDoCmd(wq)
}()
return ch
}
// Shutdown waits all worker to finish
func (w *Worker) Shutdown() {
w.stopOnce.Do(func() {
@ -91,15 +110,15 @@ func (w *Worker) loop() {
func (w *Worker) workDoCmd(req workRequest) {
var rt Response
if len(req.Cmd) == 1 {
rt = w.workDoSingle(req.Cmd[0])
rt = w.workDoSingle(req.Context, req.Cmd[0])
} else {
rt = w.workDoGroup(req.Cmd, req.PipeMapping)
rt = w.workDoGroup(req.Context, req.Cmd, req.PipeMapping)
}
rt.RequestID = req.RequestID
req.resultCh <- rt
}
func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
func (w *Worker) workDoSingle(ctx context.Context, rc Cmd) (rt Response) {
c, copyOutSet, err := w.prepareCmd(rc)
if err != nil {
rt.Error = err
@ -109,7 +128,7 @@ func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
EnvironmentPool: w.envPool,
Cmd: c,
}
result, err := s.Run()
result, err := s.Run(ctx)
if err != nil {
rt.Error = err
return
@ -119,7 +138,7 @@ func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
return
}
func (w *Worker) workDoGroup(rc []Cmd, pm []PipeMap) (rt Response) {
func (w *Worker) workDoGroup(ctx context.Context, rc []Cmd, pm []PipeMap) (rt Response) {
var rts []Result
p := preparePipeMapping(pm)
cs := make([]*envexec.Cmd, 0, len(rc))
@ -139,7 +158,7 @@ func (w *Worker) workDoGroup(rc []Cmd, pm []PipeMap) (rt Response) {
Cmd: cs,
Pipes: p,
}
results, err := g.Run()
results, err := g.Run(ctx)
if err != nil {
rt.Error = err
return