make ConfigDir configurable (#61)

make ConfigDir configurable

* internal: retrieve ConfigDir from GOPS_CONFIG_DIR environment variable

* make code graceful

* code optimize

README.md

@@ -45,6 +45,9 @@ func main() {
 }
 ```
 
+Otherwise, you could set `GOPS_CONFIG_DIR` environment variables to assign your config dir.
+Default, gops will use the current user's home directory(AppData on windows).
+
 ### Manual
 
 It is possible to use gops tool both in local and remote mode.

agent/agent.go

@@ -44,6 +44,11 @@ type Options struct {
 	// Optional.
 	Addr string
 
+	// ConfigDir is the directory to store the configuration file,
+	// PID of the gops process, filename, port as well as content.
+	// Optional.
+	ConfigDir string
+
 	// ShutdownCleanup automatically cleans up resources if the
 	// running process receives an interrupt. Otherwise, users
 	// can call Close before shutting down.
@@ -68,11 +73,17 @@ func Listen(opts Options) error {
 		return fmt.Errorf("gops: agent already listening at: %v", listener.Addr())
 	}
 
-	gopsdir, err := internal.ConfigDir()
-	if err != nil {
-		return err
+	// new
+	gopsdir := opts.ConfigDir
+	if gopsdir == "" {
+		cfgDir, err := internal.ConfigDir()
+		if err != nil {
+			return err
+		}
+		gopsdir = cfgDir
 	}
-	err = os.MkdirAll(gopsdir, os.ModePerm)
+
+	err := os.MkdirAll(gopsdir, os.ModePerm)
 	if err != nil {
 		return err
 	}

agent/agent_test.go

@@ -32,6 +32,17 @@ func TestAgentClose(t *testing.T) {
 	}
 }
 
+func TestUseCustomConfigDir(t *testing.T) {
+	err := Listen(Options{
+		ConfigDir:       os.TempDir(),
+		ShutdownCleanup: true,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	Close()
+}
+
 func TestAgentListenMultipleClose(t *testing.T) {
 	err := Listen(Options{})
 	if err != nil {

internal/internal.go

@@ -15,7 +15,13 @@ import (
 	"strings"
 )
 
+const gopsConfigDirEnvKey = "GOPS_CONFIG_DIR"
+
 func ConfigDir() (string, error) {
+	if configDir := os.Getenv(gopsConfigDirEnvKey); configDir != "" {
+		return configDir, nil
+	}
+
 	if runtime.GOOS == "windows" {
 		return filepath.Join(os.Getenv("APPDATA"), "gops"), nil
 	}

internal/internal_test.go

@@ -0,0 +1,23 @@
+package internal
+
+import (
+	"os"
+	"testing"
+)
+
+func TestConfigDir(t *testing.T) {
+	key := gopsConfigDirEnvKey
+	oldDir := os.Getenv(key)
+	defer os.Setenv(key, oldDir)
+
+	newDir := "foo-bar"
+	os.Setenv(key, newDir)
+	configDir, err := ConfigDir()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if g, w := configDir, newDir; g != w {
+		t.Errorf("ConfigDir: got=%v want=%v", g, w)
+	}
+}