Update heap dump capture for compressed file to use the original extension

Author: ycrash

internal/capture/heap.go

@@ -43,6 +43,25 @@ var compressedHeapExtensions = []string{
 const hdOut = "heap_dump.out"
 const hdZip = "heap_dump.zip"
 
+type HeapDumpFile struct {
+	SrcHeapDumpPath string
+	SrcHeapDumpFile *os.File
+
+	DstCompressedPath string
+	DstCompressedFile *os.File
+}
+
+func (h *HeapDumpFile) IsSrcHeapDumpAlreadyCompressed() bool {
+	ext := strings.TrimPrefix(filepath.Ext(h.SrcHeapDumpPath), ".")
+
+	for _, compressedExt := range compressedHeapExtensions {
+		if ext == compressedExt {
+			return true
+		}
+	}
+	return false
+}
+
 type HeapDump struct {
 	Capture
 	JavaHome string
@@ -61,84 +80,120 @@ func NewHeapDump(javaHome string, pid int, hdPath string, dump bool) *HeapDump {
 	}
 }
 
-// getDestHeapDumpFilename returns the appropriate filename based on compression status
-func (t *HeapDump) getDestHeapDumpFilename(isCompressed bool) string {
-	if isCompressed {
-		return hdZip
-	}
-	return hdOut
-}
-
-// Run executes the heap dump capture process and uploads the captured file
-// to the specified endpoint.
 func (t *HeapDump) Run() (Result, error) {
-	var hd *os.File
-	var err error
-	var isCompressed bool
-	var contentEncoding string
+	heapDumpFile := &HeapDumpFile{}
 
-	// Try pre-captured heap dump first
 	if len(t.hdPath) > 0 {
-		isCompressed, contentEncoding = isCompressedHeapFile(t.hdPath)
-		if isCompressed {
-			logger.Log("detected pre-compressed heap dump file: %s with encoding: %s", t.hdPath, contentEncoding)
+		heapDumpFile.SrcHeapDumpPath = t.hdPath
+
+		hd, err := os.Open(heapDumpFile.SrcHeapDumpPath)
+
+		// Fallback, try to open the file in the Docker container
+		if err != nil && runtime.GOOS == "linux" {
+			logger.Log("failed to open hdPath(%s) err: %s. Trying to open in the Docker container...", t.hdPath, err.Error())
+			hd, err = os.Open(filepath.Join("/proc", strconv.Itoa(t.Pid), "root", t.hdPath))
 		}
 
-		hd, err = t.getPreCapturedDumpFile(isCompressed)
 		if err != nil {
+			logger.Log("failed to open hdPath(%s) err: %s", t.hdPath, err.Error())
 			return Result{
-				Msg: fmt.Sprintf("capture heap dump failed: %s", err.Error()),
+				Msg: fmt.Sprintf("failed to open heap dump file: %s", err.Error()),
 				Ok:  false,
-			}, nil
+			}, err
 		}
+
+		heapDumpFile.SrcHeapDumpFile = hd
+
+		// Ensure the source heap dump file is closed when the function exits
+		defer func() {
+			if err := hd.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
+				logger.Log("failed to close source heap dump file: %s", err.Error())
+			}
+		}()
 	} else if t.Pid > 0 && t.dump {
-		var actualDumpPath string
-		// Then try capturing a new heap dump
-		hd, actualDumpPath, err = t.captureDumpFile()
+		hd, actualDumpPath, err := t.captureDumpFile()
 		if err != nil {
 			return Result{
 				Msg: fmt.Sprintf("capture heap dump failed: %s", err.Error()),
 				Ok:  false,
 			}, nil
 		}
 
-		// Cleanup the actual dump file after we're done with it
-		if actualDumpPath != "" {
-			defer func() {
-				err := os.Remove(actualDumpPath)
-				if err != nil {
-					logger.Trace().Err(err).Str("file", actualDumpPath).Msg("failed to rm hd file")
-				}
-			}()
-		}
+		heapDumpFile.SrcHeapDumpPath = actualDumpPath
+		heapDumpFile.SrcHeapDumpFile = hd
+
+		// Ensure the captured heap dump file is closed when the function exits
+		defer func() {
+			if err := hd.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
+				logger.Log("failed to close captured heap dump file: %s", err.Error())
+			}
+		}()
+
+		// because this code creates the file, it's responsible for cleaning it up
+		defer func() {
+			if removeErr := os.Remove(actualDumpPath); removeErr != nil {
+				logger.Log("failed to rm hd file %s cause err: %s", actualDumpPath, removeErr.Error())
+			}
+		}()
 	}
 
-	if hd == nil {
+	if heapDumpFile.SrcHeapDumpFile == nil {
 		return Result{Msg: "skipped heap dump"}, nil
 	}
 
-	tempFilename := t.getDestHeapDumpFilename(isCompressed)
+	if heapDumpFile.IsSrcHeapDumpAlreadyCompressed() {
+		logger.Log("copying heap dump data %s", t.hdPath)
+
+		srcExtension := strings.TrimPrefix(filepath.Ext(heapDumpFile.SrcHeapDumpPath), ".")
+		heapDumpFile.DstCompressedPath = "heap_dump." + srcExtension
 
-	defer func() {
-		err := hd.Close()
-		if err != nil && !errors.Is(err, os.ErrClosed) {
-			logger.Log("failed to close hd file %s cause err: %s", tempFilename, err.Error())
+		dstCompressedFile, err := os.OpenFile(heapDumpFile.DstCompressedPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
+		if err != nil {
+			return Result{
+				Msg: fmt.Sprintf("failed creating heap dump in current working directory: %s", err.Error()),
+				Ok:  false,
+			}, nil
 		}
-	}()
 
-	var fileToUpload *os.File
-	var uploadContentEncoding string
+		// Ensure the file is closed when the function exits
+		defer func() {
+			if err := dstCompressedFile.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
+				logger.Log("failed to close destination compressed file: %s", err.Error())
+			}
+		}()
+
+		_, err = io.Copy(dstCompressedFile, heapDumpFile.SrcHeapDumpFile)
+		if err != nil {
+			return Result{
+				Msg: fmt.Sprintf("failed copying heap dump data: %s", err.Error()),
+				Ok:  false,
+			}, nil
+		}
 
-	if isCompressed {
-		// If the file is already compressed, use it directly without re-compressing
-		logger.Log("file is already compressed, skipping compression step")
-		fileToUpload = hd
-		uploadContentEncoding = contentEncoding
+		// Sync the file to ensure all data is written to disk
+		err = dstCompressedFile.Sync()
+		if err != nil {
+			return Result{
+				Msg: fmt.Sprintf("failed syncing heap dump file: %s", err.Error()),
+				Ok:  false,
+			}, nil
+		}
+
+		// Reset file position to beginning for reading during upload
+		_, err = dstCompressedFile.Seek(0, 0)
+		if err != nil {
+			return Result{
+				Msg: fmt.Sprintf("failed seeking to beginning of heap dump file: %s", err.Error()),
+				Ok:  false,
+			}, nil
+		}
+
+		heapDumpFile.DstCompressedFile = dstCompressedFile
+		logger.Log("copied heap dump data %s to %s", t.hdPath, heapDumpFile.DstCompressedPath)
 	} else {
-		// For uncompressed files, compress them
 		logger.Log("captured heap dump data, zipping...")
 
-		zipfile, err := t.CreateZipFile(hd)
+		zipfile, err := t.CreateZipFile(heapDumpFile.SrcHeapDumpFile)
 		if err != nil {
 			return Result{
 				Msg: fmt.Sprintf("capture heap dump failed: %s", err.Error()),
@@ -152,65 +207,16 @@ func (t *HeapDump) Run() (Result, error) {
 			}
 		}()
 
-		defer func() {
-			err = os.Remove(tempFilename)
-			if err != nil {
-				logger.Log("failed to rm hd file %s cause err: %s", tempFilename, err.Error())
-			}
-		}()
-
-		fileToUpload = zipfile
-		uploadContentEncoding = "zip"
+		heapDumpFile.DstCompressedPath = hdZip
+		heapDumpFile.DstCompressedFile = zipfile
 	}
 
-	result := t.UploadCapturedFile(fileToUpload, uploadContentEncoding)
+	// Upload heapDumpFile.DstCompressedFile to the endpoint
+	dstCompressedExtension := strings.TrimPrefix(filepath.Ext(heapDumpFile.DstCompressedPath), ".")
+	result := t.UploadCapturedFile(heapDumpFile.DstCompressedFile, dstCompressedExtension)
 	return result, nil
 }
 
-// getPreCapturedDumpFile handles the case when a heap dump is pre-captured (using the hdPath field)
-func (t *HeapDump) getPreCapturedDumpFile(isCompressed bool) (*os.File, error) {
-	hdf, err := os.Open(t.hdPath)
-
-	// Fallback, try to open the file in the Docker container
-	if err != nil && runtime.GOOS == "linux" {
-		logger.Log("failed to open hdPath(%s) err: %s. Trying to open in the Docker container...", t.hdPath, err.Error())
-		hdf, err = os.Open(filepath.Join("/proc", strconv.Itoa(t.Pid), "root", t.hdPath))
-	}
-
-	if err != nil {
-		logger.Log("failed to open hdPath(%s) err: %s", t.hdPath, err.Error())
-		return nil, err
-	}
-
-	logger.Log("copying heap dump data %s", t.hdPath)
-
-	defer func() {
-		err := hdf.Close()
-		if err != nil {
-			logger.Log("failed to close hd file %s cause err: %s", t.hdPath, err.Error())
-		}
-	}()
-
-	filename := t.getDestHeapDumpFilename(isCompressed)
-	hd, err := os.Create(filename)
-	if err != nil {
-		return nil, err
-	}
-
-	_, err = io.Copy(hd, hdf)
-	if err != nil {
-		return nil, err
-	}
-
-	_, err = hd.Seek(0, 0)
-	if err != nil {
-		return nil, err
-	}
-
-	logger.Log("copied heap dump data %s to %s", t.hdPath, filename)
-	return hd, nil
-}
-
 // captureDumpFile handles the case when a heap dump needs to be captured (using the Pid field)
 // and returns both the file handle and the actual dump path
 func (t *HeapDump) captureDumpFile() (*os.File, string, error) {
@@ -255,25 +261,44 @@ func (t *HeapDump) CreateZipFile(hd *os.File) (*os.File, error) {
 		return nil, fmt.Errorf("failed to create zip file: %w", err)
 	}
 
-	writer := zip.NewWriter(bufio.NewWriter(zipfile))
+	bufferedWriter := bufio.NewWriter(zipfile)
+	writer := zip.NewWriter(bufferedWriter)
 	out, err := writer.Create(hdOut)
 	if err != nil {
+		zipfile.Close()
 		return nil, fmt.Errorf("failed to create zip file: %w", err)
 	}
 
 	_, err = io.Copy(out, hd)
 	if err != nil {
+		zipfile.Close()
 		return nil, fmt.Errorf("failed to zip heap dump file: %w", err)
 	}
 
+	// Close zip writer first to write central directory records
 	err = writer.Close()
 	if err != nil {
-		return nil, fmt.Errorf("failed to finish zipping heap dump file: %w", err)
+		zipfile.Close()
+		return nil, fmt.Errorf("failed to close zip writer: %w", err)
+	}
+
+	// Flush the buffered writer to ensure all data is written to the file
+	err = bufferedWriter.Flush()
+	if err != nil {
+		zipfile.Close()
+		return nil, fmt.Errorf("failed to flush zip file buffer: %w", err)
+	}
+
+	// Close the file to ensure all data is synced to disk
+	err = zipfile.Close()
+	if err != nil {
+		return nil, fmt.Errorf("failed to close zip file: %w", err)
 	}
 
-	e := zipfile.Sync()
-	if e != nil && !errors.Is(e, os.ErrClosed) {
-		logger.Log("failed to sync file %s", e)
+	// Reopen the file for reading to pass to the upload function
+	zipfile, err = os.Open(hdZip)
+	if err != nil {
+		return nil, fmt.Errorf("failed to reopen zip file for upload: %w", err)
 	}
 
 	return zipfile, nil
@@ -288,19 +313,6 @@ func (t *HeapDump) UploadCapturedFile(file *os.File, contentEncoding string) Res
 	}
 }
 
-// isCompressedHeapFile checks if the given file is a compressed heap dump file
-// and returns the extension as the MIME type if it is.
-func isCompressedHeapFile(filePath string) (bool, string) {
-	ext := strings.TrimPrefix(filepath.Ext(filePath), ".")
-
-	for _, compressedExt := range compressedHeapExtensions {
-		if ext == compressedExt {
-			return true, ext
-		}
-	}
-	return false, ext
-}
-
 // heapDump runs the JDK tool (jcmd, jattach, etc) to capture the heap dump to the requested file.
 // The returned actualDumpPath is the actual file name written to is returned.
 // In IBM JDK, this may not be the same as the requested filename for several reasons: