浏览代码

Merge pull request #45 from rauls5382/nodir

More careful handling of temporary directories
Josef Jelinek 8 年前
父节点
当前提交
8b3c6b12a7
共有 2 个文件被更改,包括 14 次插入15 次删除
  1. 1
    1
      internal/driver/commands.go
  2. 13
    14
      internal/driver/fetch.go

+ 1
- 1
internal/driver/commands.go 查看文件

411
 			return err
411
 			return err
412
 		}
412
 		}
413
 
413
 
414
-		tempFile, err := newTempFile(os.Getenv("PPROF_TMPDIR"), "pprof", "."+suffix)
414
+		tempFile, err := newTempFile(os.TempDir(), "pprof", "."+suffix)
415
 		if err != nil {
415
 		if err != nil {
416
 			return err
416
 			return err
417
 		}
417
 		}

+ 13
- 14
internal/driver/fetch.go 查看文件

37
 // there are some failures. It will return an error if it is unable to
37
 // there are some failures. It will return an error if it is unable to
38
 // fetch any profiles.
38
 // fetch any profiles.
39
 func fetchProfiles(s *source, o *plugin.Options) (*profile.Profile, error) {
39
 func fetchProfiles(s *source, o *plugin.Options) (*profile.Profile, error) {
40
-	if err := setTmpDir(o.UI); err != nil {
41
-		return nil, err
42
-	}
43
-
44
 	sources := make([]profileSource, 0, len(s.Sources)+len(s.Base))
40
 	sources := make([]profileSource, 0, len(s.Sources)+len(s.Base))
45
 	for _, src := range s.Sources {
41
 	for _, src := range s.Sources {
46
 		sources = append(sources, profileSource{
42
 		sources = append(sources, profileSource{
76
 
72
 
77
 	// Save a copy of the merged profile if there is at least one remote source.
73
 	// Save a copy of the merged profile if there is at least one remote source.
78
 	if save {
74
 	if save {
75
+		dir, err := setTmpDir(o.UI)
76
+		if err != nil {
77
+			return nil, err
78
+		}
79
+
79
 		prefix := "pprof."
80
 		prefix := "pprof."
80
 		if len(p.Mapping) > 0 && p.Mapping[0].File != "" {
81
 		if len(p.Mapping) > 0 && p.Mapping[0].File != "" {
81
 			prefix += filepath.Base(p.Mapping[0].File) + "."
82
 			prefix += filepath.Base(p.Mapping[0].File) + "."
84
 			prefix += s.Type + "."
85
 			prefix += s.Type + "."
85
 		}
86
 		}
86
 
87
 
87
-		dir := os.Getenv("PPROF_TMPDIR")
88
 		tempFile, err := newTempFile(dir, prefix, ".pb.gz")
88
 		tempFile, err := newTempFile(dir, prefix, ".pb.gz")
89
 		if err == nil {
89
 		if err == nil {
90
 			if err = p.Write(tempFile); err == nil {
90
 			if err = p.Write(tempFile); err == nil {
211
 	err    error
211
 	err    error
212
 }
212
 }
213
 
213
 
214
-// setTmpDir sets the PPROF_TMPDIR environment variable with a new
215
-// temp directory, if not already set.
216
-func setTmpDir(ui plugin.UI) error {
214
+// setTmpDir prepares the directory to use to save profiles retrieved
215
+// remotely. It is selected from PPROF_TMPDIR, defaults to $HOME/pprof.
216
+func setTmpDir(ui plugin.UI) (string, error) {
217
 	if profileDir := os.Getenv("PPROF_TMPDIR"); profileDir != "" {
217
 	if profileDir := os.Getenv("PPROF_TMPDIR"); profileDir != "" {
218
-		return nil
218
+		return profileDir, nil
219
 	}
219
 	}
220
-	for _, tmpDir := range []string{os.Getenv("HOME") + "/pprof", "/tmp"} {
220
+	for _, tmpDir := range []string{os.Getenv("HOME") + "/pprof", os.TempDir()} {
221
 		if err := os.MkdirAll(tmpDir, 0755); err != nil {
221
 		if err := os.MkdirAll(tmpDir, 0755); err != nil {
222
 			ui.PrintErr("Could not use temp dir ", tmpDir, ": ", err.Error())
222
 			ui.PrintErr("Could not use temp dir ", tmpDir, ": ", err.Error())
223
 			continue
223
 			continue
224
 		}
224
 		}
225
-		os.Setenv("PPROF_TMPDIR", tmpDir)
226
-		return nil
225
+		return tmpDir, nil
227
 	}
226
 	}
228
-	return fmt.Errorf("failed to identify temp dir")
227
+	return "", fmt.Errorf("failed to identify temp dir")
229
 }
228
 }
230
 
229
 
231
 // grabProfile fetches a profile. Returns the profile, sources for the
230
 // grabProfile fetches a profile. Returns the profile, sources for the
424
 	ui.Print(fmt.Sprintf(
423
 	ui.Print(fmt.Sprintf(
425
 		"Converting %s to a profile.proto... (May take a few minutes)",
424
 		"Converting %s to a profile.proto... (May take a few minutes)",
426
 		perfPath))
425
 		perfPath))
427
-	profile, err := newTempFile("/tmp", "pprof_", ".pb.gz")
426
+	profile, err := newTempFile(os.TempDir(), "pprof_", ".pb.gz")
428
 	if err != nil {
427
 	if err != nil {
429
 		return nil, err
428
 		return nil, err
430
 	}
429
 	}