Implement Cancel() API for cancelling builds

Author: EmmEff

client/build.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 
 	jsonresp "github.com/sylabs/json-resp"
@@ -42,3 +43,23 @@ func (c *Client) Submit(ctx context.Context, br BuildRequest) (bi BuildInfo, err
 	}
 	return
 }
+
+// Cancel cancels an existing build. The context controls the lifetime of the
+// request.
+func (c *Client) Cancel(ctx context.Context, buildID string) error {
+	req, err := c.newRequest(http.MethodPut, fmt.Sprintf("/v1/build/%s/_cancel", buildID), "", nil)
+	if err != nil {
+		return err
+	}
+	c.Logger.Logf("Sending build cancellation request to %s", req.URL.String())
+
+	res, err := c.HTTPClient.Do(req.WithContext(ctx))
+	if err != nil {
+		return err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusNoContent {
+		return fmt.Errorf("build cancellation failed: http status %d", res.StatusCode)
+	}
+	return nil
+}

client/build_test.go

@@ -87,3 +87,29 @@ func TestSubmit(t *testing.T) {
 		})
 	}
 }
+
+func TestCancel(t *testing.T) {
+	// Start a mock server
+	m := mockService{t: t}
+	s := httptest.NewServer(&m)
+	defer s.Close()
+
+	// Enough of a struct to test with
+	url, err := url.Parse(s.URL)
+	if err != nil {
+		t.Fatalf("failed to parse URL: %v", err)
+	}
+	c, err := client.New(&client.Config{
+		BaseURL: url.String(),
+	})
+	if err != nil {
+		t.Fatalf("failed to parse URL: %v", err)
+	}
+
+	m.cancelResponseCode = 204
+
+	err = c.Cancel(context.Background(), "00000000")
+	if err != nil {
+		t.Fatal(err)
+	}
+}

client/client_test.go

@@ -29,6 +29,7 @@ type mockService struct {
 	wsCloseCode        int
 	statusResponseCode int
 	imageResponseCode  int
+	cancelResponseCode int
 	httpAddr           string
 }
 
@@ -39,12 +40,13 @@ func TestMain(m *testing.M) {
 }
 
 const (
-	authToken      = "auth_token"
-	stdoutContents = "some_output"
-	imageContents  = "image_contents"
-	buildPath      = "/v1/build"
-	wsPath         = "/v1/build-ws/"
-	imagePath      = "/v1/image"
+	authToken         = "auth_token"
+	stdoutContents    = "some_output"
+	imageContents     = "image_contents"
+	buildPath         = "/v1/build"
+	wsPath            = "/v1/build-ws/"
+	imagePath         = "/v1/image"
+	buildCancelSuffix = "/_cancel"
 )
 
 func newResponse(m *mockService, id string, def []byte, libraryRef string) client.BuildInfo {
@@ -110,6 +112,15 @@ func (m *mockService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				m.t.Fatal(err)
 			}
 		}
+	} else if r.Method == http.MethodPut && strings.HasSuffix(r.RequestURI, buildCancelSuffix) {
+		// Mock build cancellation endpoint
+		if m.cancelResponseCode == http.StatusNoContent {
+			w.WriteHeader(http.StatusNoContent)
+		} else {
+			if err := jsonresp.WriteError(w, "", m.cancelResponseCode); err != nil {
+				m.t.Fatal(err)
+			}
+		}
 	} else {
 		w.WriteHeader(http.StatusNotFound)
 	}

client/output.go

@@ -10,6 +10,9 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/gorilla/websocket"
 )
@@ -37,21 +40,31 @@ func (c *Client) GetOutput(ctx context.Context, buildID string, or OutputReader)
 	h := http.Header{}
 	c.setRequestHeaders(h)
 
-	ws, resp, err := websocket.DefaultDialer.Dial(u.String(), h)
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	ws, resp, err := websocket.DefaultDialer.DialContext(ctx, u.String(), h)
 	if err != nil {
 		c.Logger.Logf("websocket dial err - %s, partial response: %+v", err, resp)
 		return err
 	}
 	defer ws.Close()
 
-	for {
-		// Check if context has expired
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
+	go func() {
+		sigCh := make(chan os.Signal, 1)
+		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+
+		fmt.Printf("\rShutting down due to signal: %v\n", <-sigCh)
+
+		if err := c.Cancel(ctx, buildID); err != nil {
+			c.Logger.Logf("build cancellation request failed: %v", err)
 		}
 
+		cancel()
+
+	}()
+
+	for {
 		// Read from websocket
 		mt, msg, err := ws.ReadMessage()
 		if err != nil {