transport: context package support allowing cancellation of any network operation

plumbing/transport/common.go

@@ -13,6 +13,7 @@
 package transport
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -72,7 +73,7 @@ type UploadPackSession interface {
 	// including a packfile. Don't be confused by terminology, the client
 	// side of a git-upload-pack is called git-fetch-pack, although here
 	// the same interface is used to make it RPC-like.
-	UploadPack(*packp.UploadPackRequest) (*packp.UploadPackResponse, error)
+	UploadPack(context.Context, *packp.UploadPackRequest) (*packp.UploadPackResponse, error)
 }
 
 // ReceivePackSession represents a git-receive-pack session.
@@ -86,7 +87,7 @@ type ReceivePackSession interface {
 	// terminology, the client side of a git-receive-pack is called
 	// git-send-pack, although here the same interface is used to make it
 	// RPC-like.
-	ReceivePack(*packp.ReferenceUpdateRequest) (*packp.ReportStatus, error)
+	ReceivePack(context.Context, *packp.ReferenceUpdateRequest) (*packp.ReportStatus, error)
 }
 
 // Endpoint represents a Git URL in any supported protocol.

plumbing/transport/file/client.go

@@ -30,7 +30,9 @@ func NewClient(uploadPackBin, receivePackBin string) transport.Transport {
 	})
 }
 
-func (r *runner) Command(cmd string, ep transport.Endpoint, auth transport.AuthMethod) (common.Command, error) {
+func (r *runner) Command(cmd string, ep transport.Endpoint, auth transport.AuthMethod,
+) (common.Command, error) {
+
 	switch cmd {
 	case transport.UploadPackServiceName:
 		cmd = r.UploadPackBin
@@ -72,6 +74,11 @@ func (c *command) StdoutPipe() (io.Reader, error) {
 	return c.cmd.StdoutPipe()
 }
 
+func (c *command) Kill() error {
+	c.cmd.Process.Kill()
+	return c.Close()
+}
+
 // Close waits for the command to exit.
 func (c *command) Close() error {
 	if c.closed {
@@ -81,6 +88,7 @@ func (c *command) Close() error {
 	defer func() {
 		c.closed = true
 		_ = c.stderrCloser.Close()
+
 	}()
 
 	err := c.cmd.Wait()

plumbing/transport/file/upload_pack_test.go

@@ -78,3 +78,9 @@ func (s *UploadPackSuite) TestNonExistentCommand(c *C) {
 	c.Assert(err, ErrorMatches, ".*file.*")
 	c.Assert(session, IsNil)
 }
+
+func (s *UploadPackSuite) TestUploadPackWithContextOnRead(c *C) {
+	// TODO: Fix race condition when Session.Close and the read failed due to a
+	// canceled context when the packfile is being read.
+	c.Skip("UploadPack has a race condition when we Close the session")
+}

plumbing/transport/http/receive_pack.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -25,7 +26,7 @@ func (s *rpSession) AdvertisedReferences() (*packp.AdvRefs, error) {
 	return advertisedReferences(s.session, transport.ReceivePackServiceName)
 }
 
-func (s *rpSession) ReceivePack(req *packp.ReferenceUpdateRequest) (
+func (s *rpSession) ReceivePack(ctx context.Context, req *packp.ReferenceUpdateRequest) (
 	*packp.ReportStatus, error) {
 	url := fmt.Sprintf(
 		"%s/%s",
@@ -37,7 +38,7 @@ func (s *rpSession) ReceivePack(req *packp.ReferenceUpdateRequest) (
 		return nil, err
 	}
 
-	res, err := s.doRequest(http.MethodPost, url, buf)
+	res, err := s.doRequest(ctx, http.MethodPost, url, buf)
 	if err != nil {
 		return nil, err
 	}
@@ -61,7 +62,10 @@ func (s *rpSession) ReceivePack(req *packp.ReferenceUpdateRequest) (
 	return report, report.Error()
 }
 
-func (s *rpSession) doRequest(method, url string, content *bytes.Buffer) (*http.Response, error) {
+func (s *rpSession) doRequest(
+	ctx context.Context, method, url string, content *bytes.Buffer,
+) (*http.Response, error) {
+
 	var body io.Reader
 	if content != nil {
 		body = content
@@ -75,7 +79,7 @@ func (s *rpSession) doRequest(method, url string, content *bytes.Buffer) (*http.
 	applyHeadersToRequest(req, content, s.endpoint.Host(), transport.ReceivePackServiceName)
 	s.applyAuthToRequest(req)
 
-	res, err := s.client.Do(req)
+	res, err := s.client.Do(req.WithContext(ctx))
 	if err != nil {
 		return nil, plumbing.NewUnexpectedError(err)
 	}

plumbing/transport/http/upload_pack.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -28,7 +29,10 @@ func (s *upSession) AdvertisedReferences() (*packp.AdvRefs, error) {
 	return advertisedReferences(s.session, transport.UploadPackServiceName)
 }
 
-func (s *upSession) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackResponse, error) {
+func (s *upSession) UploadPack(
+	ctx context.Context, req *packp.UploadPackRequest,
+) (*packp.UploadPackResponse, error) {
+
 	if req.IsEmpty() {
 		return nil, transport.ErrEmptyUploadPackRequest
 	}
@@ -47,7 +51,7 @@ func (s *upSession) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackR
 		return nil, err
 	}
 
-	res, err := s.doRequest(http.MethodPost, url, content)
+	res, err := s.doRequest(ctx, http.MethodPost, url, content)
 	if err != nil {
 		return nil, err
 	}
@@ -70,7 +74,10 @@ func (s *upSession) Close() error {
 	return nil
 }
 
-func (s *upSession) doRequest(method, url string, content *bytes.Buffer) (*http.Response, error) {
+func (s *upSession) doRequest(
+	ctx context.Context, method, url string, content *bytes.Buffer,
+) (*http.Response, error) {
+
 	var body io.Reader
 	if content != nil {
 		body = content
@@ -84,7 +91,7 @@ func (s *upSession) doRequest(method, url string, content *bytes.Buffer) (*http.
 	applyHeadersToRequest(req, content, s.endpoint.Host(), transport.UploadPackServiceName)
 	s.applyAuthToRequest(req)
 
-	res, err := s.client.Do(req)
+	res, err := s.client.Do(req.WithContext(ctx))
 	if err != nil {
 		return nil, plumbing.NewUnexpectedError(err)
 	}

plumbing/transport/internal/common/common.go

@@ -7,6 +7,7 @@ package common
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -64,6 +65,13 @@ type Command interface {
 	Close() error
 }
 
+// CommandKiller expands the Command interface, enableing it for being killed.
+type CommandKiller interface {
+	// Kill and close the session whatever the state it is. It will block until
+	// the command is terminated.
+	Kill() error
+}
+
 type client struct {
 	cmdr Commander
 }
@@ -212,7 +220,7 @@ func (s *session) handleAdvRefDecodeError(err error) error {
 
 // UploadPack performs a request to the server to fetch a packfile. A reader is
 // returned with the packfile content. The reader must be closed after reading.
-func (s *session) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackResponse, error) {
+func (s *session) UploadPack(ctx context.Context, req *packp.UploadPackRequest) (*packp.UploadPackResponse, error) {
 	if req.IsEmpty() {
 		return nil, transport.ErrEmptyUploadPackRequest
 	}
@@ -227,11 +235,14 @@ func (s *session) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackRes
 
 	s.packRun = true
 
-	if err := uploadPack(s.Stdin, s.Stdout, req); err != nil {
+	in := s.StdinContext(ctx)
+	out := s.StdoutContext(ctx)
+
+	if err := uploadPack(in, out, req); err != nil {
 		return nil, err
 	}
 
-	r, err := ioutil.NonEmptyReader(s.Stdout)
+	r, err := ioutil.NonEmptyReader(out)
 	if err == ioutil.ErrEmptyReader {
 		if c, ok := s.Stdout.(io.Closer); ok {
 			_ = c.Close()
@@ -244,22 +255,45 @@ func (s *session) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackRes
 		return nil, err
 	}
 
-	rc := ioutil.NewReadCloser(r, s.Command)
+	rc := ioutil.NewReadCloser(r, s)
 	return DecodeUploadPackResponse(rc, req)
 }
 
-func (s *session) ReceivePack(req *packp.ReferenceUpdateRequest) (*packp.ReportStatus, error) {
+func (s *session) StdinContext(ctx context.Context) io.WriteCloser {
+	return ioutil.NewWriteCloserOnError(
+		ioutil.NewContextWriteCloser(ctx, s.Stdin),
+		s.onError,
+	)
+}
+
+func (s *session) StdoutContext(ctx context.Context) io.Reader {
+	return ioutil.NewReaderOnError(
+		ioutil.NewContextReader(ctx, s.Stdout),
+		s.onError,
+	)
+}
+
+func (s *session) onError(err error) {
+	if k, ok := s.Command.(CommandKiller); ok {
+		_ = k.Kill()
+	}
+
+	_ = s.Close()
+}
+
+func (s *session) ReceivePack(ctx context.Context, req *packp.ReferenceUpdateRequest) (*packp.ReportStatus, error) {
 	if _, err := s.AdvertisedReferences(); err != nil {
 		return nil, err
 	}
 
 	s.packRun = true
 
-	if err := req.Encode(s.Stdin); err != nil {
+	w := s.StdinContext(ctx)
+	if err := req.Encode(w); err != nil {
 		return nil, err
 	}
 
-	if err := s.Stdin.Close(); err != nil {
+	if err := w.Close(); err != nil {
 		return nil, err
 	}
 
@@ -270,11 +304,12 @@ func (s *session) ReceivePack(req *packp.ReferenceUpdateRequest) (*packp.ReportS
 	}
 
 	report := packp.NewReportStatus()
-	if err := report.Decode(s.Stdout); err != nil {
+	if err := report.Decode(s.StdoutContext(ctx)); err != nil {
 		return nil, err
 	}
 
 	if err := report.Error(); err != nil {
+		defer s.Close()
 		return report, err
 	}
 
@@ -300,8 +335,9 @@ func (s *session) finish() error {
 }
 
 func (s *session) Close() (err error) {
-	defer ioutil.CheckClose(s.Command, &err)
 	err = s.finish()
+
+	defer ioutil.CheckClose(s.Command, &err)
 	return
 }
 

plumbing/transport/internal/common/server.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -34,7 +35,7 @@ func ServeUploadPack(cmd ServerCommand, s transport.UploadPackSession) (err erro
 	}
 
 	var resp *packp.UploadPackResponse
-	resp, err = s.UploadPack(req)
+	resp, err = s.UploadPack(context.TODO(), req)
 	if err != nil {
 		return err
 	}
@@ -57,7 +58,7 @@ func ServeReceivePack(cmd ServerCommand, s transport.ReceivePackSession) error {
 		return fmt.Errorf("error decoding: %s", err)
 	}
 
-	rs, err := s.ReceivePack(req)
+	rs, err := s.ReceivePack(context.TODO(), req)
 	if rs != nil {
 		if err := rs.Encode(cmd.Stdout); err != nil {
 			return fmt.Errorf("error in encoding report status %s", err)

plumbing/transport/server/server.go

@@ -3,6 +3,7 @@
 package server
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -14,6 +15,7 @@ import (
 	"gopkg.in/src-d/go-git.v4/plumbing/revlist"
 	"gopkg.in/src-d/go-git.v4/plumbing/storer"
 	"gopkg.in/src-d/go-git.v4/plumbing/transport"
+	"gopkg.in/src-d/go-git.v4/utils/ioutil"
 )
 
 var DefaultServer = NewServer(DefaultLoader)
@@ -129,7 +131,7 @@ func (s *upSession) AdvertisedReferences() (*packp.AdvRefs, error) {
 	return ar, nil
 }
 
-func (s *upSession) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackResponse, error) {
+func (s *upSession) UploadPack(ctx context.Context, req *packp.UploadPackRequest) (*packp.UploadPackResponse, error) {
 	if req.IsEmpty() {
 		return nil, transport.ErrEmptyUploadPackRequest
 	}
@@ -167,7 +169,9 @@ func (s *upSession) UploadPack(req *packp.UploadPackRequest) (*packp.UploadPackR
 		pw.CloseWithError(err)
 	}()
 
-	return packp.NewUploadPackResponseWithPackfile(req, pr), nil
+	return packp.NewUploadPackResponseWithPackfile(req,
+		ioutil.NewContextReadCloser(ctx, pr),
+	), nil
 }
 
 func (s *upSession) objectsToUpload(req *packp.UploadPackRequest) ([]plumbing.Hash, error) {
@@ -222,7 +226,7 @@ var (
 	ErrUpdateReference = errors.New("failed to update ref")
 )
 
-func (s *rpSession) ReceivePack(req *packp.ReferenceUpdateRequest) (*packp.ReportStatus, error) {
+func (s *rpSession) ReceivePack(ctx context.Context, req *packp.ReferenceUpdateRequest) (*packp.ReportStatus, error) {
 	if s.caps == nil {
 		s.caps = capability.NewList()
 		if err := s.setSupportedCapabilities(s.caps); err != nil {
@@ -238,7 +242,8 @@ func (s *rpSession) ReceivePack(req *packp.ReferenceUpdateRequest) (*packp.Repor
 
 	//TODO: Implement 'atomic' update of references.
 
-	if err := s.writePackfile(req.Packfile); err != nil {
+	r := ioutil.NewContextReadCloser(ctx, req.Packfile)
+	if err := s.writePackfile(r); err != nil {
 		s.unpackErr = err
 		s.firstErr = err
 		return s.reportStatus(), err

plumbing/transport/server/upload_pack_test.go

@@ -39,6 +39,10 @@ func (s *UploadPackSuite) TestAdvertisedReferencesNotExists(c *C) {
 	c.Assert(r, IsNil)
 }
 
+func (s *UploadPackSuite) TestUploadPackWithContext(c *C) {
+	c.Skip("UploadPack cannot be canceled on server")
+}
+
 // Tests server with `asClient = true`. This is recommended when using a server
 // registered directly with `client.InstallProtocol`.
 type ClientLikeUploadPackSuite struct {