Merge branch 'stream-helpers' into 'master'

Use gitaly-proto/helper stream io helpers

See merge request !133

5 files changed, 78 insertions, 83 deletions

diff --git a/internal/service/smarthttp/inforefs.go b/internal/service/smarthttp/inforefs.go
index 0ea48d529..59026f7a5 100644
--- a/internal/service/smarthttp/inforefs.go
+++ b/internal/service/smarthttp/inforefs.go
@@ -6,34 +6,25 @@ import (
 	"log"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go"
+	pbhelper "gitlab.com/gitlab-org/gitaly-proto/go/helper"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 )
 
-type infoRefsResponseWriter struct {
-	infoRefsResponseServer
-}
-
-type infoRefsResponseServer interface {
-	Send(*pb.InfoRefsResponse) error
-}
-
-func (w infoRefsResponseWriter) Write(p []byte) (int, error) {
-	if err := w.Send(&pb.InfoRefsResponse{Data: p}); err != nil {
-		return 0, err
-	}
-
-	return len(p), nil
-}
-
 func (s *server) InfoRefsUploadPack(in *pb.InfoRefsRequest, stream pb.SmartHTTP_InfoRefsUploadPackServer) error {
-	return handleInfoRefs("upload-pack", in.Repository, infoRefsResponseWriter{stream})
+	w := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.InfoRefsResponse{Data: p})
+	})
+	return handleInfoRefs("upload-pack", in.Repository, w)
 }
 
 func (s *server) InfoRefsReceivePack(in *pb.InfoRefsRequest, stream pb.SmartHTTP_InfoRefsReceivePackServer) error {
-	return handleInfoRefs("receive-pack", in.Repository, infoRefsResponseWriter{stream})
+	w := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.InfoRefsResponse{Data: p})
+	})
+	return handleInfoRefs("receive-pack", in.Repository, w)
 }
 
 func handleInfoRefs(service string, repo *pb.Repository, w io.Writer) error {
diff --git a/internal/service/smarthttp/receive_pack.go b/internal/service/smarthttp/receive_pack.go
index bde99adf6..6c4ddfefe 100644
--- a/internal/service/smarthttp/receive_pack.go
+++ b/internal/service/smarthttp/receive_pack.go
@@ -8,18 +8,11 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go"
+	pbhelper "gitlab.com/gitlab-org/gitaly-proto/go/helper"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 )
 
-type receivePackBytesReader struct {
-	pb.SmartHTTP_PostReceivePackServer
-}
-
-type receivePackWriter struct {
-	pb.SmartHTTP_PostReceivePackServer
-}
-
 func (s *server) PostReceivePack(stream pb.SmartHTTP_PostReceivePackServer) error {
 	req, err := stream.Recv() // First request contains only Repository and GlId
 	if err != nil {
@@ -29,9 +22,13 @@ func (s *server) PostReceivePack(stream pb.SmartHTTP_PostReceivePackServer) erro
 		return err
 	}
 
-	streamBytesReader := receivePackBytesReader{stream}
-	stdin := &streamReader{br: streamBytesReader}
-	stdout := receivePackWriter{stream}
+	stdin := pbhelper.NewReceiveReader(func() ([]byte, error) {
+		resp, err := stream.Recv()
+		return resp.GetData(), err
+	})
+	stdout := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.PostReceivePackResponse{Data: p})
+	})
 	glIDEnv := fmt.Sprintf("GL_ID=%s", req.GlId)
 	repoPath, err := helper.GetRepoPath(req.Repository)
 	if err != nil {
@@ -65,20 +62,3 @@ func validateReceivePackRequest(req *pb.PostReceivePackRequest) error {
 
 	return nil
 }
-
-func (rw receivePackWriter) Write(p []byte) (int, error) {
-	resp := &pb.PostReceivePackResponse{Data: p}
-	if err := rw.Send(resp); err != nil {
-		return 0, err
-	}
-	return len(p), nil
-}
-
-func (br receivePackBytesReader) ReceiveBytes() ([]byte, error) {
-	resp, err := br.Recv()
-	if err != nil {
-		return nil, err
-	}
-
-	return resp.GetData(), nil
-}
diff --git a/internal/service/smarthttp/upload_pack.go b/internal/service/smarthttp/upload_pack.go
index 9528155ee..f34ee1b7f 100644
--- a/internal/service/smarthttp/upload_pack.go
+++ b/internal/service/smarthttp/upload_pack.go
@@ -7,18 +7,11 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go"
+	pbhelper "gitlab.com/gitlab-org/gitaly-proto/go/helper"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 )
 
-type uploadPackBytesReader struct {
-	pb.SmartHTTP_PostUploadPackServer
-}
-
-type uploadPackWriter struct {
-	pb.SmartHTTP_PostUploadPackServer
-}
-
 func (s *server) PostUploadPack(stream pb.SmartHTTP_PostUploadPackServer) error {
 	req, err := stream.Recv() // First request contains Repository only
 	if err != nil {
@@ -28,9 +21,13 @@ func (s *server) PostUploadPack(stream pb.SmartHTTP_PostUploadPackServer) error
 		return err
 	}
 
-	streamBytesReader := uploadPackBytesReader{stream}
-	stdin := &streamReader{br: streamBytesReader}
-	stdout := uploadPackWriter{stream}
+	stdin := pbhelper.NewReceiveReader(func() ([]byte, error) {
+		resp, err := stream.Recv()
+		return resp.GetData(), err
+	})
+	stdout := pbhelper.NewSendWriter(func(p []byte) error {
+		return stream.Send(&pb.PostUploadPackResponse{Data: p})
+	})
 	repoPath, err := helper.GetRepoPath(req.Repository)
 	if err != nil {
 		return err
@@ -60,20 +57,3 @@ func validateUploadPackRequest(req *pb.PostUploadPackRequest) error {
 
 	return nil
 }
-
-func (rw uploadPackWriter) Write(p []byte) (int, error) {
-	resp := &pb.PostUploadPackResponse{Data: p}
-	if err := rw.Send(resp); err != nil {
-		return 0, err
-	}
-	return len(p), nil
-}
-
-func (br uploadPackBytesReader) ReceiveBytes() ([]byte, error) {
-	resp, err := br.Recv()
-	if err != nil {
-		return nil, err
-	}
-
-	return resp.GetData(), nil
-}
diff --git a/vendor/gitlab.com/gitlab-org/gitaly-proto/go/helper/stream.go b/vendor/gitlab.com/gitlab-org/gitaly-proto/go/helper/stream.go
new file mode 100644
index 000000000..77cd157e9
--- /dev/null
+++ b/vendor/gitlab.com/gitlab-org/gitaly-proto/go/helper/stream.go
@@ -0,0 +1,44 @@
+package helper
+
+import (
+	"io"
+)
+
+// NewReceiveReader turns receiver into an io.Reader. Errors from the
+// receiver function are passed on unmodified. This means receiver should
+// emit io.EOF when done.
+func NewReceiveReader(receiver func() ([]byte, error)) io.Reader {
+	return &receiveReader{receiver: receiver}
+}
+
+type receiveReader struct {
+	receiver func() ([]byte, error)
+	data     []byte
+	err      error
+}
+
+func (rr *receiveReader) Read(p []byte) (int, error) {
+	if len(rr.data) == 0 {
+		rr.data, rr.err = rr.receiver()
+	}
+	n := copy(p, rr.data)
+	rr.data = rr.data[n:]
+	if len(rr.data) == 0 {
+		return n, rr.err
+	}
+	return n, nil
+}
+
+// NewSendWriter turns sender into an io.Writer. The number of 'bytes
+// written' reported back is always len(p).
+func NewSendWriter(sender func(p []byte) error) io.Writer {
+	return &sendWriter{sender: sender}
+}
+
+type sendWriter struct {
+	sender func([]byte) error
+}
+
+func (sw *sendWriter) Write(p []byte) (int, error) {
+	return len(p), sw.sender(p)
+}
diff --git a/vendor/vendor.json b/vendor/vendor.json
index 824df9f73..785662085 100644
--- a/vendor/vendor.json
+++ b/vendor/vendor.json
@@ -119,18 +119,18 @@
 		{
 			"checksumSHA1": "qnITGZYkVMWqbOeVVt6jMKJ544M=",
 			"path": "gitlab.com/gitlab-org/gitaly-proto/go",
-			"revision": "52f77b23166e640a932e50223472d761404afb42",
-			"revisionTime": "2017-03-29T16:52:58Z",
-			"version": "v0.5.0",
-			"versionExact": "v0.5.0"
+			"revision": "8aafca712c95b248e3d6f3b5f5d60043595f94c5",
+			"revisionTime": "2017-04-25T14:57:47Z",
+			"version": "reader-writer-helpers",
+			"versionExact": "reader-writer-helpers"
 		},
 		{
-			"checksumSHA1": "2E36lBoyaVky2EJP1E5ub6Rg+uI=",
+			"checksumSHA1": "GkeSZfXVbtAkBZOrswot19GJZqQ=",
 			"path": "gitlab.com/gitlab-org/gitaly-proto/go/helper",
-			"revision": "607fd598bc02ecc750d905ea8ba233673ff0a86b",
-			"revisionTime": "2017-03-28T10:04:06Z",
-			"version": "v0.4.0",
-			"versionExact": "v0.4.0"
+			"revision": "8aafca712c95b248e3d6f3b5f5d60043595f94c5",
+			"revisionTime": "2017-04-25T14:57:47Z",
+			"version": "reader-writer-helpers",
+			"versionExact": "reader-writer-helpers"
 		},
 		{
 			"checksumSHA1": "9jjO5GjLa0XF/nfWihF02RoH4qc=",