From 48dbb593f7cc0dceb55d232ac198e82f3df1c964 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B8rn=20Erik=20Pedersen?=
 <bjorn.erik.pedersen@gmail.com>
Date: Thu, 25 Jun 2020 11:44:27 +0200
Subject: [PATCH] commands: Add an option to print memory usage at intervals

Use it with `hugo --print-mem
---
 commands/commands.go |  2 ++
 commands/hugo.go     | 54 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/commands/commands.go b/commands/commands.go
index 66fd9caa..09e0c845 100644
--- a/commands/commands.go
+++ b/commands/commands.go
@@ -212,6 +212,7 @@ type hugoBuilderCommon struct {
 	memprofile   string
 	mutexprofile string
 	traceprofile string
+	printm       bool
 
 	// TODO(bep) var vs string
 	logging    bool
@@ -299,6 +300,7 @@ func (cc *hugoBuilderCommon) handleFlags(cmd *cobra.Command) {
 	cmd.Flags().BoolP("path-warnings", "", false, "print warnings on duplicate target paths etc.")
 	cmd.Flags().StringVarP(&cc.cpuprofile, "profile-cpu", "", "", "write cpu profile to `file`")
 	cmd.Flags().StringVarP(&cc.memprofile, "profile-mem", "", "", "write memory profile to `file`")
+	cmd.Flags().BoolVarP(&cc.printm, "print-mem", "", false, "print memory usage to screen at intervals")
 	cmd.Flags().StringVarP(&cc.mutexprofile, "profile-mutex", "", "", "write Mutex profile to `file`")
 	cmd.Flags().StringVarP(&cc.traceprofile, "trace", "", "", "write trace to `file` (not useful in general)")
 
diff --git a/commands/hugo.go b/commands/hugo.go
index b7392bc4..5442c32d 100644
--- a/commands/hugo.go
+++ b/commands/hugo.go
@@ -82,6 +82,7 @@ func (r Response) IsUserError() bool {
 // Execute adds all child commands to the root command HugoCmd and sets flags appropriately.
 // The args are usually filled with os.Args[1:].
 func Execute(args []string) Response {
+
 	hugoCmd := newCommandsBuilder().addAll().build()
 	cmd := hugoCmd.getCommand()
 	cmd.SetArgs(args)
@@ -427,7 +428,37 @@ func (c *commandeer) initMutexProfile() (func(), error) {
 
 }
 
+func (c *commandeer) initMemTicker() func() {
+	memticker := time.NewTicker(5 * time.Second)
+	quit := make(chan struct{})
+	printMem := func() {
+		var m runtime.MemStats
+		runtime.ReadMemStats(&m)
+		fmt.Printf("\n\nAlloc = %v\nTotalAlloc = %v\nSys = %v\nNumGC = %v\n\n", formatByteCount(m.Alloc), formatByteCount(m.TotalAlloc), formatByteCount(m.Sys), m.NumGC)
+
+	}
+
+	go func() {
+		for {
+			select {
+			case <-memticker.C:
+				printMem()
+			case <-quit:
+				memticker.Stop()
+				printMem()
+				return
+			}
+
+		}
+	}()
+
+	return func() {
+		close(quit)
+	}
+}
+
 func (c *commandeer) initProfiling() (func(), error) {
+
 	stopCPUProf, err := c.initCPUProfile()
 	if err != nil {
 		return nil, err
@@ -443,6 +474,11 @@ func (c *commandeer) initProfiling() (func(), error) {
 		return nil, err
 	}
 
+	var stopMemTicker func()
+	if c.h.printm {
+		stopMemTicker = c.initMemTicker()
+	}
+
 	return func() {
 		c.initMemProfile()
 
@@ -456,6 +492,10 @@ func (c *commandeer) initProfiling() (func(), error) {
 		if stopTraceProf != nil {
 			stopTraceProf()
 		}
+
+		if stopMemTicker != nil {
+			stopMemTicker()
+		}
 	}, nil
 }
 
@@ -1175,3 +1215,17 @@ func pickOneWriteOrCreatePath(events []fsnotify.Event) string {
 
 	return name
 }
+
+func formatByteCount(b uint64) string {
+	const unit = 1000
+	if b < unit {
+		return fmt.Sprintf("%d B", b)
+	}
+	div, exp := int64(unit), 0
+	for n := b / unit; n >= unit; n /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f %cB",
+		float64(b)/float64(div), "kMGTPE"[exp])
+}