From 81f564835eca343ad4a154057894dcc280a445a0 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Thu, 11 Jul 2024 19:39:28 -0700
Subject: [PATCH] `profiletool`: Use better compression ratio when writing
 profiles.

The Go profile library only writes compressed profiles using the "best
speed" (i.e. largest size) compression ratio setting, and does not have
a way to set it. This change avoids calling the Go profile library's
default compression code and writes it uncompressed, piped to a writer
with the same compression algorithm but higher compression ratio.

This is useful because this tool is meant to be used to process profiles
before they are checked into the repository (so that they are available at
build time for PGO builds). Since they are checked in the repo, we need to
minimize their size to keep the repository size small.

PiperOrigin-RevId: 651613043
---
 tools/profiletool/profiletool.go | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/tools/profiletool/profiletool.go b/tools/profiletool/profiletool.go
index a7ae15910..6f0d8a736 100644
--- a/tools/profiletool/profiletool.go
+++ b/tools/profiletool/profiletool.go
@@ -16,6 +16,7 @@
 package main
 
 import (
+	"compress/gzip"
 	"errors"
 	"fmt"
 	"os"
@@ -109,14 +110,11 @@ func mergeProfiles() error {
 	if err != nil {
 		return fmt.Errorf("cannot create %q: %w", *mergeOut, err)
 	}
-	if err := merged.Write(mergedFile); err != nil {
-		mergedFile.Close()
+	defer mergedFile.Close()
+	if err := writeMaxCompressionProfile(merged, mergedFile); err != nil {
 		os.Remove(*mergeOut)
 		return fmt.Errorf("cannot write merged profile to %q: %w", *mergeOut, err)
 	}
-	if err := mergedFile.Close(); err != nil {
-		return fmt.Errorf("cannot close %q: %w", *mergeOut, err)
-	}
 	return nil
 }
 
@@ -143,7 +141,7 @@ func compactProfile() error {
 	if err != nil {
 		return fmt.Errorf("cannot create %q: %w", *compactOut, err)
 	}
-	if err := prof.Write(compactedFile); err != nil {
+	if err := writeMaxCompressionProfile(prof, compactedFile); err != nil {
 		compactedFile.Close()
 		os.Remove(*compactOut)
 		return fmt.Errorf("cannot write compacted profile to %q: %w", *compactOut, err)
@@ -154,6 +152,25 @@ func compactProfile() error {
 	return nil
 }
 
+// writeMaxCompressionProfile writes a profile to a file with the maximum
+// compression level. The file handle is not closed.
+func writeMaxCompressionProfile(p *profile.Profile, out *os.File) error {
+	// The profile library writes profiles with the fastest (i.e. worst)
+	// compression level by default, and does not allow setting the compression
+	// level. So we compress it with the maximum level manually here.
+	writer, err := gzip.NewWriterLevel(out, gzip.BestCompression)
+	if err != nil {
+		return fmt.Errorf("cannot create zlib writer: %w", err)
+	}
+	if err := p.WriteUncompressed(writer); err != nil {
+		return fmt.Errorf("cannot write profile to zlib writer: %w", err)
+	}
+	if err := writer.Close(); err != nil {
+		return fmt.Errorf("cannot close zlib writer: %w", err)
+	}
+	return nil
+}
+
 func main() {
 	if len(os.Args) < 2 {
 		printUsage()