fix: return serialize errors instead of panicking (task 383)

2 files changed, 43 insertions, 9 deletions

diff --git a/internal/flamegraph/iordata.go b/internal/flamegraph/iordata.go
index 13b78fd..a205916 100644
--- a/internal/flamegraph/iordata.go
+++ b/internal/flamegraph/iordata.go
@@ -3,6 +3,7 @@ package flamegraph
 import (
 	"bytes"
 	"encoding/gob"
+	"errors"
 	"fmt"
 	"io"
 	"iter"
@@ -26,6 +27,8 @@ type tidType = uint32
 type flagsType = file.Flags
 type pathMap map[pathType]map[traceIdType]map[commType]map[pidType]map[tidType]map[flagsType]Counter
 
+var hostnameFn = os.Hostname
+
 type recordKey struct {
 	Path    pathType
 	TraceID traceIdType
@@ -98,10 +101,10 @@ func (iod iorData) merge(other iorData) iorData {
 	return iod
 }
 
-func (iod iorData) serializeToFile(flamegraphName string) error {
-	hostname, err := os.Hostname()
+func (iod iorData) serializeToFile(flamegraphName string) (retErr error) {
+	hostname, err := hostnameFn()
 	if err != nil {
-		panic(err)
+		return fmt.Errorf("get hostname: %w", err)
 	}
 	if flamegraphName == "" {
 		flamegraphName = "default"
@@ -114,22 +117,33 @@ func (iod iorData) serializeToFile(flamegraphName string) error {
 
 	file, err := os.Create(tmpFilename)
 	if err != nil {
-		return err
+		return fmt.Errorf("create temp file %s: %w", tmpFilename, err)
 	}
-	defer file.Close()
+	defer func() {
+		if err := file.Close(); err != nil {
+			retErr = errors.Join(retErr, fmt.Errorf("close temp file %s: %w", tmpFilename, err))
+		}
+	}()
 
 	encoder := zstd.NewWriter(file)
-	defer encoder.Close()
+	defer func() {
+		if err := encoder.Close(); err != nil {
+			retErr = errors.Join(retErr, fmt.Errorf("close zstd writer for %s: %w", tmpFilename, err))
+		}
+	}()
 
 	gobEncoder := gob.NewEncoder(encoder)
 	if err := gobEncoder.Encode(iod.records); err != nil {
-		return err
+		return fmt.Errorf("encode ior records: %w", err)
 	}
 	if err := encoder.Flush(); err != nil {
-		return err
+		return fmt.Errorf("flush ior records: %w", err)
 	}
 
-	return os.Rename(tmpFilename, filename)
+	if err := os.Rename(tmpFilename, filename); err != nil {
+		return fmt.Errorf("rename %s to %s: %w", tmpFilename, filename, err)
+	}
+	return nil
 }
 
 func (iod *iorData) loadFromFile(filename string) error {
diff --git a/internal/flamegraph/iordata_test.go b/internal/flamegraph/iordata_test.go
index 5e95976..f4855b0 100644
--- a/internal/flamegraph/iordata_test.go
+++ b/internal/flamegraph/iordata_test.go
@@ -2,6 +2,8 @@ package flamegraph
 
 import (
 	"bytes"
+	"errors"
+	"strings"
 	"syscall"
 	"testing"
 
@@ -288,6 +290,24 @@ func TestDeserializeInvalidData(t *testing.T) {
 	}
 }
 
+func TestSerializeToFileHostnameErrorReturnsError(t *testing.T) {
+	origHostnameFn := hostnameFn
+	t.Cleanup(func() { hostnameFn = origHostnameFn })
+
+	hostnameFn = func() (string, error) {
+		return "", errors.New("hostname unavailable")
+	}
+
+	iod := newIorData()
+	err := iod.serializeToFile("test")
+	if err == nil {
+		t.Fatal("Expected error when hostname lookup fails, got nil")
+	}
+	if !strings.Contains(err.Error(), "get hostname") {
+		t.Fatalf("Expected get hostname context, got %v", err)
+	}
+}
+
 func bothArraysHaveSameElements(a, b []string) bool {
 	if len(a) != len(b) {
 		return false