From f5360b0074185238519d5e7367a24fed82c95d94 Mon Sep 17 00:00:00 2001
From: Vaibhav Kamra <vkamra@alcion.ai>
Date: Tue, 29 Nov 2022 21:54:47 -0800
Subject: [PATCH] Memory profiling

---
 build/Dockerfile |  4 +++-
 src/cli/cli.go   | 27 +++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/build/Dockerfile b/build/Dockerfile
index 209024254..6925c6f66 100644
--- a/build/Dockerfile
+++ b/build/Dockerfile
@@ -19,6 +19,8 @@ COPY --from=builder /go/src/app/corso /corso
 #  Pull tls certs directly from latest upstream image
 COPY --from=alpine:latest /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/
 
+RUN apk add --no-cache --update binutils util-linux logrotate
+
 ENV CORSO_HOME=/app/corso
 ENV CORSO_CONFIG_DIR=$CORSO_HOME \
 	KOPIA_CONFIG_PATH=$CORSO_HOME/kopia/config/repository.config \
@@ -28,4 +30,4 @@ ENV CORSO_CONFIG_DIR=$CORSO_HOME \
 	KOPIA_PERSIST_CREDENTIALS_ON_CONNECT=false \
 	KOPIA_CHECK_FOR_UPDATES=false
 
-ENTRYPOINT ["/corso"]
\ No newline at end of file
+ENTRYPOINT ["tail", "-f", "/dev/null"]
diff --git a/src/cli/cli.go b/src/cli/cli.go
index 378675003..bb1e93bbc 100644
--- a/src/cli/cli.go
+++ b/src/cli/cli.go
@@ -2,8 +2,11 @@ package cli
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"regexp"
+	"runtime"
+	"runtime/pprof"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -95,6 +98,16 @@ func Handle() {
 	ctx, log := logger.Seed(ctx, logger.PreloadLogLevel())
 	defer func() {
 		_ = log.Sync() // flush all logs in the buffer
+		f, err := os.Create("mem.prof")
+		if err != nil {
+			log.Fatal("could not create memory profile: ", err)
+		}
+		defer f.Close() // error handling omitted for example
+		runtime.GC()    // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("could not write memory profile: ", err)
+		}
+		PrintMemUsage()
 	}()
 
 	if err := corsoCmd.ExecuteContext(ctx); err != nil {
@@ -113,3 +126,17 @@ func indentExamplesTemplate(template string) string {
 
 	return e.ReplaceAllString(template, "{{.Example | indent 2}}")
 }
+
+func PrintMemUsage() {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
+	fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
+	fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
+	fmt.Printf("\tNumGC = %v\n", m.NumGC)
+}
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}