From e7bae9df81f9499b09352161aaaa90ea3dcea19c Mon Sep 17 00:00:00 2001
From: Dan Winship <danwinship@redhat.com>
Date: Sat, 19 Feb 2022 11:27:57 -0500
Subject: [PATCH] Count iptables lines as we write them

---
 pkg/proxy/iptables/proxier.go |  6 ++----
 pkg/proxy/util/utils.go       | 20 ++++++++++++++------
 pkg/proxy/util/utils_test.go  |  5 ++++-
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/pkg/proxy/iptables/proxier.go b/pkg/proxy/iptables/proxier.go
index 0855096c1b1..d0af939db76 100644
--- a/pkg/proxy/iptables/proxier.go
+++ b/pkg/proxy/iptables/proxier.go
@@ -1507,10 +1507,8 @@ func (proxier *Proxier) syncProxyRules() {
 		"-j", "ACCEPT",
 	)
 
-	numberFilterIptablesRules := utilproxy.CountBytesLines(proxier.filterRules.Bytes())
-	metrics.IptablesRulesTotal.WithLabelValues(string(utiliptables.TableFilter)).Set(float64(numberFilterIptablesRules))
-	numberNatIptablesRules := utilproxy.CountBytesLines(proxier.natRules.Bytes())
-	metrics.IptablesRulesTotal.WithLabelValues(string(utiliptables.TableNAT)).Set(float64(numberNatIptablesRules))
+	metrics.IptablesRulesTotal.WithLabelValues(string(utiliptables.TableFilter)).Set(float64(proxier.filterRules.Lines()))
+	metrics.IptablesRulesTotal.WithLabelValues(string(utiliptables.TableNAT)).Set(float64(proxier.natRules.Lines()))
 
 	// Write the end-of-table markers.
 	proxier.filterRules.Write("COMMIT")
diff --git a/pkg/proxy/util/utils.go b/pkg/proxy/util/utils.go
index f20fc194307..1cd78a4a8ca 100644
--- a/pkg/proxy/util/utils.go
+++ b/pkg/proxy/util/utils.go
@@ -521,7 +521,8 @@ func GetClusterIPByFamily(ipFamily v1.IPFamily, service *v1.Service) string {
 }
 
 type LineBuffer struct {
-	b bytes.Buffer
+	b     bytes.Buffer
+	lines int
 }
 
 // Write takes a list of arguments, each a string or []string, joins all the
@@ -547,22 +548,34 @@ func (buf *LineBuffer) Write(args ...interface{}) {
 		}
 	}
 	buf.b.WriteByte('\n')
+	buf.lines++
 }
 
 // WriteBytes writes bytes to buffer, and terminates with newline.
 func (buf *LineBuffer) WriteBytes(bytes []byte) {
 	buf.b.Write(bytes)
 	buf.b.WriteByte('\n')
+	buf.lines++
 }
 
+// Reset clears buf
 func (buf *LineBuffer) Reset() {
 	buf.b.Reset()
+	buf.lines = 0
 }
 
+// Bytes returns the contents of buf as a []byte
 func (buf *LineBuffer) Bytes() []byte {
 	return buf.b.Bytes()
 }
 
+// Lines returns the number of lines in buf. Note that more precisely, this returns the
+// number of times Write() or WriteBytes() was called; it assumes that you never wrote
+// any newlines to the buffer yourself.
+func (buf *LineBuffer) Lines() int {
+	return buf.lines
+}
+
 // RevertPorts is closing ports in replacementPortsMap but not in originalPortsMap. In other words, it only
 // closes the ports opened in this sync.
 func RevertPorts(replacementPortsMap, originalPortsMap map[netutils.LocalPort]netutils.Closeable) {
@@ -574,8 +587,3 @@ func RevertPorts(replacementPortsMap, originalPortsMap map[netutils.LocalPort]ne
 		}
 	}
 }
-
-// CountBytesLines counts the number of lines in a bytes slice
-func CountBytesLines(b []byte) int {
-	return bytes.Count(b, []byte{'\n'})
-}
diff --git a/pkg/proxy/util/utils_test.go b/pkg/proxy/util/utils_test.go
index 4b0c83e1a87..ed5ae7b4d4f 100644
--- a/pkg/proxy/util/utils_test.go
+++ b/pkg/proxy/util/utils_test.go
@@ -1185,6 +1185,9 @@ func TestLineBufferWrite(t *testing.T) {
 			if want, got := testCase.expected, string(testBuffer.Bytes()); !strings.EqualFold(want, got) {
 				t.Fatalf("write word is %v\n expected: %q, got: %q", testCase.input, want, got)
 			}
+			if testBuffer.Lines() != 1 {
+				t.Fatalf("expected 1 line, got: %d", testBuffer.Lines())
+			}
 		})
 	}
 }
@@ -1267,7 +1270,7 @@ func TestWriteCountLines(t *testing.T) {
 			for i := 0; i < testCase.expected; i++ {
 				testBuffer.Write(randSeq())
 			}
-			n := CountBytesLines(testBuffer.Bytes())
+			n := testBuffer.Lines()
 			if n != testCase.expected {
 				t.Fatalf("lines expected: %d, got: %d", testCase.expected, n)
 			}