Use stream to send debug files (ethereum#279)

* Fix protobuf file * Small fix * Use grpc streams to send big debug files
streamingfast · Jan 13, 2022 · 4829114 · 4829114
1 parent 3b2bfa0
commit 4829114
Show file tree

Hide file tree

Showing 8 changed files with 433 additions and 119 deletions.
diff --git a/command/debug.go b/command/debug.go
@@ -6,7 +6,6 @@ import (
 	"archive/tar"
 	"compress/gzip"
 	"context"
-	"encoding/hex"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -19,6 +18,10 @@ import (
 
 	"github.com/ethereum/go-ethereum/command/flagset"
 	"github.com/ethereum/go-ethereum/command/server/proto"
+	"github.com/golang/protobuf/jsonpb"
+	gproto "github.com/golang/protobuf/proto"
+	"github.com/golang/protobuf/ptypes/empty"
+	grpc_net_conn "github.com/mitchellh/go-grpc-net-conn"
 )
 
 type DebugCommand struct {
@@ -118,16 +121,35 @@ func (d *DebugCommand) Run(args []string) int {
 			req.Type = proto.PprofRequest_LOOKUP
 			req.Profile = profile
 		}
-		resp, err := clt.Pprof(ctx, req)
+		stream, err := clt.Pprof(ctx, req)
 		if err != nil {
 			return err
 		}
-		// write file
-		raw, err := hex.DecodeString(resp.Payload)
+		// wait for open request
+		msg, err := stream.Recv()
 		if err != nil {
 			return err
 		}
-		if err := ioutil.WriteFile(filepath.Join(tmp, filename+".prof"), raw, 0755); err != nil {
+		if _, ok := msg.Event.(*proto.PprofResponse_Open_); !ok {
+			return fmt.Errorf("expected open message")
+		}
+
+		// create the stream
+		conn := &grpc_net_conn.Conn{
+			Stream:   stream,
+			Response: &proto.PprofResponse_Input{},
+			Decode: grpc_net_conn.SimpleDecoder(func(msg gproto.Message) *[]byte {
+				return &msg.(*proto.PprofResponse_Input).Data
+			}),
+		}
+
+		file, err := os.OpenFile(filepath.Join(tmp, filename+".prof"), os.O_RDWR|os.O_CREATE, 0644)
+		if err != nil {
+			return err
+		}
+		defer file.Close()
+
+		if _, err := io.Copy(file, conn); err != nil {
 			return err
 		}
 		return nil
@@ -148,6 +170,25 @@ func (d *DebugCommand) Run(args []string) int {
 		}
 	}
 
+	// append the status
+	{
+		statusResp, err := clt.Status(ctx, &empty.Empty{})
+		if err != nil {
+			d.UI.Output(fmt.Sprintf("Failed to get status: %v", err))
+			return 1
+		}
+		m := jsonpb.Marshaler{}
+		data, err := m.MarshalToString(statusResp)
+		if err != nil {
+			d.UI.Output(err.Error())
+			return 1
+		}
+		if err := ioutil.WriteFile(filepath.Join(tmp, "status.json"), []byte(data), 0644); err != nil {
+			d.UI.Output(fmt.Sprintf("Failed to write status: %v", err))
+			return 1
+		}
+	}
+
 	// Exit before archive if output directory was specified
 	if d.output != "" {
 		d.UI.Output(fmt.Sprintf("Created debug directory: %s", tmp))

diff --git a/command/server/proto/server.pb.go b/command/server/proto/server.pb.go
diff --git a/command/server/proto/server.proto b/command/server/proto/server.proto
@@ -7,7 +7,7 @@ import "google/protobuf/empty.proto";
 option go_package = "/command/server/proto";
 
 service Bor {
-    rpc Pprof(PprofRequest) returns (PprofResponse);
+    rpc Pprof(PprofRequest) returns (stream PprofResponse);
 
     rpc PeersAdd(PeersAddRequest) returns (PeersAddResponse);
 
@@ -112,6 +112,18 @@ message PprofRequest {
 }
 
 message PprofResponse {
-    string payload = 1;
-    map<string, string> headers = 2;
+    oneof event {
+        Open open = 1;
+        Input input = 2;
+        google.protobuf.Empty eof = 3;
+    }
+
+    message Open {
+        map<string, string> headers = 1;
+        int64 size = 2;
+    }
+
+    message Input {
+        bytes data = 1;    
+    }
 }
diff --git a/command/server/proto/server_grpc.pb.go b/command/server/proto/server_grpc.pb.go
diff --git a/command/server/service.go b/command/server/service.go
@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"encoding/hex"
 	"fmt"
 	"math/big"
 	"reflect"
@@ -13,14 +12,17 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	gproto "github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/ptypes/empty"
+	grpc_net_conn "github.com/mitchellh/go-grpc-net-conn"
 )
 
-func (s *Server) Pprof(ctx context.Context, req *proto.PprofRequest) (*proto.PprofResponse, error) {
+func (s *Server) Pprof(req *proto.PprofRequest, stream proto.Bor_PprofServer) error {
 	var payload []byte
 	var headers map[string]string
 	var err error
 
+	ctx := context.Background()
 	switch req.Type {
 	case proto.PprofRequest_CPU:
 		payload, headers, err = pprof.CPUProfile(ctx, int(req.Seconds))
@@ -30,14 +32,42 @@ func (s *Server) Pprof(ctx context.Context, req *proto.PprofRequest) (*proto.Ppr
 		payload, headers, err = pprof.Profile(req.Profile, 0, 0)
 	}
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	resp := &proto.PprofResponse{
-		Payload: hex.EncodeToString(payload),
-		Headers: headers,
+	// open the stream and send the headers
+	err = stream.Send(&proto.PprofResponse{
+		Event: &proto.PprofResponse_Open_{
+			Open: &proto.PprofResponse_Open{
+				Headers: headers,
+				Size:    int64(len(payload)),
+			},
+		},
+	})
+	if err != nil {
+		return err
 	}
-	return resp, nil
+
+	// Wrap our conn around the response.
+	conn := &grpc_net_conn.Conn{
+		Stream:  stream,
+		Request: &proto.PprofResponse_Input{},
+		Encode: grpc_net_conn.SimpleEncoder(func(msg gproto.Message) *[]byte {
+			return &msg.(*proto.PprofResponse_Input).Data
+		}),
+	}
+	if _, err := conn.Write(payload); err != nil {
+		return err
+	}
+
+	// send the eof
+	err = stream.Send(&proto.PprofResponse{
+		Event: &proto.PprofResponse_Eof{},
+	})
+	if err != nil {
+		return err
+	}
+	return nil
 }
 
 func (s *Server) PeersAdd(ctx context.Context, req *proto.PeersAddRequest) (*proto.PeersAddResponse, error) {

diff --git a/go.mod b/go.mod
@@ -53,6 +53,7 @@ require (
 	github.com/mattn/go-colorable v0.1.8
 	github.com/mattn/go-isatty v0.0.12
 	github.com/mitchellh/cli v1.1.2
+	github.com/mitchellh/go-grpc-net-conn v0.0.0-20200427190222-eb030e4876f0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/naoina/go-stringutil v0.1.0 // indirect
 	github.com/naoina/toml v0.1.2-0.20170918210437-9fafd6967416

diff --git a/go.sum b/go.sum
@@ -375,6 +375,8 @@ github.com/mitchellh/cli v1.1.2 h1:PvH+lL2B7IQ101xQL63Of8yFS2y+aDlsFcsqNc+u/Kw=
 github.com/mitchellh/cli v1.1.2/go.mod h1:6iaV0fGdElS6dPBx0EApTxHrcWvmJphyh2n8YBLPPZ4=
 github.com/mitchellh/copystructure v1.0.0 h1:Laisrj+bAB6b/yJwB5Bt3ITZhGJdqmxquMKeZ+mmkFQ=
 github.com/mitchellh/copystructure v1.0.0/go.mod h1:SNtv71yrdKgLRyLFxmLdkAbkKEFWgYaq1OVrnRcwhnw=
+github.com/mitchellh/go-grpc-net-conn v0.0.0-20200427190222-eb030e4876f0 h1:oZuel4h7224ILBLg2SlTxdaMYXDyqcVfL4Cg1PJQHZs=
+github.com/mitchellh/go-grpc-net-conn v0.0.0-20200427190222-eb030e4876f0/go.mod h1:ZCzL0JMR6qfm7VrDC8HGwVtPA8D2Ijc/edUSBw58x94=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 h1:DpOJ2HYzCv8LZP15IdmG+YdwD2luVPHITV96TkirNBM=
@@ -732,6 +734,7 @@ google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyac
 google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY=
 google.golang.org/grpc v1.26.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
 google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
+google.golang.org/grpc v1.28.1/go.mod h1:rpkK4SK4GF4Ach/+MFLZUBavHOvF2JJB5uozKKal+60=
 google.golang.org/grpc v1.33.1/go.mod h1:fr5YgcSWrqhRRxogOsw7RzIpsmvOZ6IcH4kBYTpR3n0=
 google.golang.org/grpc v1.36.0/go.mod h1:qjiiYl8FncCW8feJPdyg3v6XW24KsRHe+dy9BAGRRjU=
 google.golang.org/grpc v1.41.0/go.mod h1:U3l9uK9J0sini8mHphKoXyaqDA/8VyGnDee1zzIUK6k=

diff --git a/scripts/tools-protobuf.sh b/scripts/tools-protobuf.sh
@@ -1,5 +1,16 @@
-curl -LO https://github.com/protocolbuffers/protobuf/releases/download/v3.12.0/protoc-3.12.0-linux-x86_64.zip
-sudo unzip protoc-3.12.0-linux-x86_64.zip -d /usr/local/bin
+#!/bin/bash
 
+# Install protobuf
+PROTOC_ZIP=protoc-3.12.0-linux-x86_64.zip
+curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v3.12.0/$PROTOC_ZIP
+sudo unzip -o $PROTOC_ZIP -d /usr/local bin/protoc
+sudo unzip -o $PROTOC_ZIP -d /usr/local 'include/*'
+rm -f $PROTOC_ZIP
+
+# Change permissions to use the binary
+sudo chmod 755 -R /usr/local/bin/protoc
+sudo chmod 755 -R /usr/local/include
+
+# Install golang extensions (DO NOT CHANGE THE VERSIONS)
 go install google.golang.org/protobuf/cmd/[email protected]
 go install google.golang.org/grpc/cmd/[email protected]