agent: add debuginfo agent command to dgraph

dgraph/cmd/debuginfo/archive.go

@@ -0,0 +1,153 @@
+/*
+ * Copyright 2019-2020 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package debuginfo
+
+import (
+	"archive/tar"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/golang/glog"
+)
+
+type tarWriter interface {
+	io.Writer
+	WriteHeader(hdr *tar.Header) error
+}
+
+type walker struct {
+	baseDir  string
+	debugDir string
+	output   tarWriter
+}
+
+// walkPath function is called for each file present within the directory
+// that walker is processing. The function operates in a best effort manner
+// and tries to archive whatever it can without throwing an error.
+func (w *walker) walkPath(path string, info os.FileInfo, err error) error {
+	if err != nil {
+		glog.Errorf("Error while walking path %s: %s", path, err)
+		return nil
+	}
+	if info == nil {
+		glog.Errorf("No file info available")
+		return nil
+	}
+
+	file, err := os.Open(path)
+	if err != nil {
+		glog.Errorf("Failed to open %s: %s", path, err)
+		return nil
+	}
+	defer file.Close()
+
+	if info.IsDir() {
+		if info.Name() == w.baseDir {
+			return nil
+		}
+		glog.Errorf("Skipping directory %s", info.Name())
+		return nil
+	}
+
+	header, err := tar.FileInfoHeader(info, info.Name())
+	if err != nil {
+		glog.Errorf("Failed to prepare file info %s: %s", info.Name(), err)
+		return nil
+	}
+
+	if w.baseDir != "" {
+		header.Name = filepath.Join(w.baseDir, strings.TrimPrefix(path, w.debugDir))
+	}
+
+	if err := w.output.WriteHeader(header); err != nil {
+		glog.Errorf("Failed to write header: %s", err)
+		return nil
+	}
+
+	_, err = io.Copy(w.output, file)
+	return err
+}
+
+// createArchive creates a gzipped tar archive for the directory provided
+// by recursively traversing in the directory.
+// The final tar is placed in the same directory with the name same to the
+// archived directory.
+func createArchive(debugDir string) (string, error) {
+	archivePath := fmt.Sprintf("%s.tar", filepath.Base(debugDir))
+	file, err := os.Create(archivePath)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	writer := tar.NewWriter(file)
+	defer writer.Close()
+
+	var baseDir string
+	if info, err := os.Stat(debugDir); os.IsNotExist(err) {
+		return "", err
+	} else if err == nil && info.IsDir() {
+		baseDir = filepath.Base(debugDir)
+	}
+
+	w := &walker{
+		baseDir:  baseDir,
+		debugDir: debugDir,
+		output:   writer,
+	}
+	return archivePath, filepath.Walk(debugDir, w.walkPath)
+}
+
+// Creates a Gzipped tar archive of the directory provided as parameter.
+func createGzipArchive(debugDir string) (string, error) {
+	source, err := createArchive(debugDir)
+	if err != nil {
+		return "", err
+	}
+
+	reader, err := os.Open(source)
+	if err != nil {
+		return "", err
+	}
+
+	filename := filepath.Base(source)
+	target := fmt.Sprintf("%s.gz", source)
+	writer, err := os.Create(target)
+	if err != nil {
+		return "", err
+	}
+	defer writer.Close()
+
+	archiver := gzip.NewWriter(writer)
+	archiver.Name = filename
+	defer archiver.Close()
+
+	_, err = io.Copy(archiver, reader)
+	if err != nil {
+		return "", err
+	}
+
+	if err = os.Remove(source); err != nil {
+		glog.Warningf("error while removing intermediate tar file: %s", err)
+	}
+
+	return target, nil
+}

dgraph/cmd/debuginfo/pprof.go

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2019-2020 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package debuginfo
+
+import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/url"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/golang/glog"
+)
+
+var pprofProfileTypes = []string{
+	"goroutine",
+	"heap",
+	"threadcreate",
+	"block",
+	"mutex",
+	"profile",
+	"trace",
+}
+
+func saveProfiles(addr, pathPrefix string, duration time.Duration, profiles []string) {
+	u, err := url.Parse(addr)
+	if err != nil || (u.Host == "" && u.Scheme != "" && u.Scheme != "file") {
+		u, err = url.Parse("http://" + addr)
+	}
+	if err != nil || u.Host == "" {
+		glog.Errorf("error while parsing address %s: %s", addr, err)
+		return
+	}
+
+	for _, profileType := range profiles {
+		source := fmt.Sprintf("%s/debug/pprof/%s?duration=%d", u.String(),
+			profileType, int(duration.Seconds()))
+		savePath := fmt.Sprintf("%s%s.gz", pathPrefix, profileType)
+
+		if err := saveProfile(source, savePath, duration); err != nil {
+			glog.Errorf("error while saving pprof profile from %s: %s", source, err)
+			continue
+		}
+
+		glog.Infof("saving %s profile in %s", profileType, savePath)
+	}
+}
+
+// saveProfile writes the profile specified in the argument fetching it from the host
+// provided in the configuration
+func saveProfile(sourceURL, filePath string, duration time.Duration) error {
+	var err error
+	var resp io.ReadCloser
+
+	glog.Infof("fetching profile over HTTP from %s", sourceURL)
+	if duration > 0 {
+		glog.Info(fmt.Sprintf("please wait... (%v)", duration))
+	}
+
+	timeout := duration + duration/2 + 2*time.Second
+	resp, err = fetchURL(sourceURL, timeout)
+	if err != nil {
+		return err
+	}
+
+	defer resp.Close()
+	out, err := os.Create(filePath)
+	if err != nil {
+		return fmt.Errorf("error while creating profile dump file: %s", err)
+	}
+	_, err = io.Copy(out, resp)
+	return err
+}
+
+// fetchURL fetches a profile from a URL using HTTP.
+func fetchURL(source string, timeout time.Duration) (io.ReadCloser, error) {
+	client := &http.Client{
+		Timeout: timeout,
+	}
+	resp, err := client.Get(source)
+	if err != nil {
+		return nil, fmt.Errorf("http fetch: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		defer resp.Body.Close()
+		return nil, statusCodeError(resp)
+	}
+
+	return resp.Body, nil
+}
+
+func statusCodeError(resp *http.Response) error {
+	if resp.Header.Get("X-Go-Pprof") != "" &&
+		strings.Contains(resp.Header.Get("Content-Type"), "text/plain") {
+		if body, err := ioutil.ReadAll(resp.Body); err == nil {
+			return fmt.Errorf("server response: %s - %s", resp.Status, body)
+		}
+	}
+	return fmt.Errorf("server response: %s", resp.Status)
+}

dgraph/cmd/debuginfo/run.go

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2019-2020 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package debuginfo
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/golang/glog"
+	"github.com/spf13/cobra"
+)
+
+type debugInfoCmdOpts struct {
+	alphaAddr string
+	zeroAddr  string
+	archive   bool
+	directory string
+	duration  uint32
+
+	pprofProfiles []string
+}
+
+var (
+	DebugInfo    x.SubCommand
+	debugInfoCmd = debugInfoCmdOpts{}
+)
+
+func init() {
+	DebugInfo.Cmd = &cobra.Command{
+		Use:   "debuginfo",
+		Short: "Generate debug info on the current node.",
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := collectDebugInfo(); err != nil {
+				glog.Errorf("error while collecting dgraph debug info: %s", err)
+				os.Exit(1)
+			}
+		},
+	}
+	DebugInfo.EnvPrefix = "DGRAPH_AGENT_DEBUGINFO"
+
+	flags := DebugInfo.Cmd.Flags()
+	flags.StringVarP(&debugInfoCmd.alphaAddr, "alpha", "a", "localhost:8080",
+		"Address of running dgraph alpha.")
+	flags.StringVarP(&debugInfoCmd.zeroAddr, "zero", "z", "", "Address of running dgraph zero.")
+	flags.StringVarP(&debugInfoCmd.directory, "directory", "d", "",
+		"Directory to write the debug info into.")
+	flags.BoolVarP(&debugInfoCmd.archive, "archive", "x", true,
+		"Whether to archive the generated report")
+	flags.Uint32VarP(&debugInfoCmd.duration, "seconds", "s", 15,
+		"Duration for time-based profile collection.")
+	flags.StringSliceVarP(&debugInfoCmd.pprofProfiles, "profiles", "p", pprofProfileTypes,
+		"List of pprof profiles to dump in the report.")
+}
+
+func collectDebugInfo() (err error) {
+	if debugInfoCmd.directory == "" {
+		debugInfoCmd.directory, err = ioutil.TempDir("/tmp", "dgraph-debuginfo")
+		if err != nil {
+			return fmt.Errorf("error while creating temporary directory: %s", err)
+		}
+	} else {
+		err = os.MkdirAll(debugInfoCmd.directory, 0644)
+		if err != nil {
+			return err
+		}
+	}
+	glog.Infof("using directory %s for debug info dump.", debugInfoCmd.directory)
+
+	collectPProfProfiles()
+
+	if debugInfoCmd.archive {
+		return archiveDebugInfo()
+	}
+	return nil
+}
+
+func collectPProfProfiles() {
+	duration := time.Duration(debugInfoCmd.duration) * time.Second
+
+	if debugInfoCmd.alphaAddr != "" {
+		filePrefix := filepath.Join(debugInfoCmd.directory, "alpha_")
+		saveProfiles(debugInfoCmd.alphaAddr, filePrefix, duration, debugInfoCmd.pprofProfiles)
+	}
+
+	if debugInfoCmd.zeroAddr != "" {
+		filePrefix := filepath.Join(debugInfoCmd.directory, "zero_")
+		saveProfiles(debugInfoCmd.zeroAddr, filePrefix, duration, debugInfoCmd.pprofProfiles)
+	}
+}
+
+func archiveDebugInfo() error {
+	archivePath, err := createGzipArchive(debugInfoCmd.directory)
+	if err != nil {
+		return fmt.Errorf("error while archiving debuginfo directory: %s", err)
+	}
+
+	glog.Infof("Debuginfo archive successful: %s", archivePath)
+
+	if err = os.RemoveAll(debugInfoCmd.directory); err != nil {
+		glog.Warningf("error while removing debuginfo directory: %s", err)
+	}
+	return nil
+}