From b33ae69800e11694c7081a06acf46b69acb3b531 Mon Sep 17 00:00:00 2001
From: Vince Prignano <vincepri@vmware.com>
Date: Fri, 5 Jul 2019 11:01:02 -0700
Subject: [PATCH] Add flag to include file's dir in logs

Signed-off-by: Vince Prignano <vincepri@vmware.com>
---
 klog.go      | 16 +++++++++++++---
 klog_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/klog.go b/klog.go
index 5e2ba8593..17d29752d 100644
--- a/klog.go
+++ b/klog.go
@@ -417,6 +417,7 @@ func InitFlags(flagset *flag.FlagSet) {
 		logging.toStderr = true
 		logging.alsoToStderr = false
 		logging.skipHeaders = false
+		logging.addDirHeader = false
 		logging.skipLogHeaders = false
 	})
 
@@ -432,6 +433,7 @@ func InitFlags(flagset *flag.FlagSet) {
 	flagset.BoolVar(&logging.toStderr, "logtostderr", logging.toStderr, "log to standard error instead of files")
 	flagset.BoolVar(&logging.alsoToStderr, "alsologtostderr", logging.alsoToStderr, "log to standard error as well as files")
 	flagset.Var(&logging.verbosity, "v", "number for the log level verbosity")
+	flagset.BoolVar(&logging.skipHeaders, "add_dir_header", logging.addDirHeader, "If true, adds the file directory to the header")
 	flagset.BoolVar(&logging.skipHeaders, "skip_headers", logging.skipHeaders, "If true, avoid header prefixes in the log messages")
 	flagset.BoolVar(&logging.skipLogHeaders, "skip_log_headers", logging.skipLogHeaders, "If true, avoid headers when opening log files")
 	flagset.Var(&logging.stderrThreshold, "stderrthreshold", "logs at or above this threshold go to stderr")
@@ -500,6 +502,9 @@ type loggingT struct {
 
 	// If true, do not add the headers to log files
 	skipLogHeaders bool
+
+	// If true, add the file directory to the header
+	addDirHeader bool
 }
 
 // buffer holds a byte Buffer for reuse. The zero value is ready for use.
@@ -585,9 +590,14 @@ func (l *loggingT) header(s severity, depth int) (*buffer, string, int) {
 		file = "???"
 		line = 1
 	} else {
-		slash := strings.LastIndex(file, "/")
-		if slash >= 0 {
-			file = file[slash+1:]
+		if slash := strings.LastIndex(file, "/"); slash >= 0 {
+			path := file
+			file = path[slash+1:]
+			if l.addDirHeader {
+				if dirsep := strings.LastIndex(path[:slash], "/"); dirsep >= 0 {
+					file = path[dirsep+1:]
+				}
+			}
 		}
 	}
 	return l.formatHeader(s, file, line), file, line
diff --git a/klog_test.go b/klog_test.go
index 22ad0bfc5..d8455aedb 100644
--- a/klog_test.go
+++ b/klog_test.go
@@ -88,6 +88,7 @@ func contains(s severity, str string, t *testing.T) bool {
 // setFlags configures the logging flags how the test expects them.
 func setFlags() {
 	logging.toStderr = false
+	logging.addDirHeader = false
 }
 
 // Test that Info works as advertised.
@@ -198,6 +199,30 @@ func TestHeader(t *testing.T) {
 	}
 }
 
+func TestHeaderWithDir(t *testing.T) {
+	setFlags()
+	logging.addDirHeader = true
+	defer logging.swap(logging.newBuffers())
+	defer func(previous func() time.Time) { timeNow = previous }(timeNow)
+	timeNow = func() time.Time {
+		return time.Date(2006, 1, 2, 15, 4, 5, .067890e9, time.Local)
+	}
+	pid = 1234
+	Info("test")
+	var line int
+	format := "I0102 15:04:05.067890    1234 klog/klog_test.go:%d] test\n"
+	n, err := fmt.Sscanf(contents(infoLog), format, &line)
+	if n != 1 || err != nil {
+		t.Errorf("log format error: %d elements, error %s:\n%s", n, err, contents(infoLog))
+	}
+	// Scanf treats multiple spaces as equivalent to a single space,
+	// so check for correct space-padding also.
+	want := fmt.Sprintf(format, line)
+	if contents(infoLog) != want {
+		t.Errorf("log format error: got:\n\t%q\nwant:\t%q", contents(infoLog), want)
+	}
+}
+
 // Test that an Error log goes to Warning and Info.
 // Even in the Info log, the source character will be E, so the data should
 // all be identical.
@@ -488,6 +513,14 @@ func BenchmarkHeader(b *testing.B) {
 	}
 }
 
+func BenchmarkHeaderWithDir(b *testing.B) {
+	logging.addDirHeader = true
+	for i := 0; i < b.N; i++ {
+		buf, _, _ := logging.header(infoLog, 0)
+		logging.putBuffer(buf)
+	}
+}
+
 // Test the logic on checking log size limitation.
 func TestFileSizeCheck(t *testing.T) {
 	setFlags()