From e5f6d2961cef719e286f3f5f7f4ab868fc4ba7cd Mon Sep 17 00:00:00 2001
From: Dmitry Vyukov <dvyukov@google.com>
Date: Fri, 5 Jul 2024 11:37:34 +0200
Subject: pkg/image: provide stats about images

---
 pkg/image/compression.go           |  9 +++++++++
 pkg/image/compression_optimized.go | 18 +++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

(limited to 'pkg/image')

diff --git a/pkg/image/compression.go b/pkg/image/compression.go
index d41392019..edc144de2 100644
--- a/pkg/image/compression.go
+++ b/pkg/image/compression.go
@@ -9,6 +9,15 @@ import (
 	"encoding/base64"
 	"fmt"
 	"io"
+	"sync/atomic"
+)
+
+var (
+	// Total amount of images in memory and consumed memory (in bytes).
+	// Currently maintained only by the optimized implementation.
+	// Cannot import stats package due to import cycles.
+	StatImages atomic.Int64
+	StatMemory atomic.Int64
 )
 
 func Compress(rawData []byte) []byte {
diff --git a/pkg/image/compression_optimized.go b/pkg/image/compression_optimized.go
index ea6be9569..819debc6e 100644
--- a/pkg/image/compression_optimized.go
+++ b/pkg/image/compression_optimized.go
@@ -22,9 +22,12 @@ type decompressScratch struct {
 	buf []byte
 }
 
+// This is just for memory consumption estimation, does not need to be precise.
+const pageSize = 4 << 10
+
 var decompressPool = sync.Pool{New: func() interface{} {
 	return &decompressScratch{
-		buf: make([]byte, 8<<10),
+		buf: make([]byte, pageSize),
 	}
 }}
 
@@ -63,11 +66,15 @@ func mustDecompress(compressed []byte) (data []byte, dtor func()) {
 	if err != nil {
 		panic(err)
 	}
+	pages := 0
 	dtor = func() {
+		StatImages.Add(-1)
+		StatMemory.Add(int64(-pages * pageSize))
 		if err := syscall.Munmap(data[:maxImageSize]); err != nil {
 			panic(err)
 		}
 	}
+	pagedIn := 0
 	offset := 0
 	for {
 		n, err := scratch.zr.Read(scratch.buf)
@@ -92,6 +99,7 @@ func mustDecompress(compressed []byte) (data []byte, dtor func()) {
 		// or whatever is the alignment for such large objects). We could also break from the middle
 		// of the loop before updating src/dst pointers, but it hurts codegen a lot (compilers like
 		// canonical loop forms).
+		hasData := false
 		words := uintptr(n-1) / wordSize
 		src := (*word)(unsafe.Pointer(&scratch.buf[0]))
 		dst := (*word)(unsafe.Pointer(&data[offset]))
@@ -101,16 +109,24 @@ func mustDecompress(compressed []byte) (data []byte, dtor func()) {
 			}
 			src = (*word)(unsafe.Pointer(uintptr(unsafe.Pointer(src)) + wordSize))
 			dst = (*word)(unsafe.Pointer(uintptr(unsafe.Pointer(dst)) + wordSize))
+			hasData = true
 		}
 		// Copy any remaining trailing bytes.
 		for i := words * wordSize; i < uintptr(n); i++ {
 			v := scratch.buf[i]
 			if v != 0 {
 				data[uintptr(offset)+i] = v
+				hasData = true
 			}
 		}
+		if hasData && offset >= pagedIn {
+			pagedIn = (offset + n + pageSize - 1) & ^(pageSize - 1)
+			pages++
+		}
 		offset += n
 	}
 	data = data[:offset]
+	StatImages.Add(1)
+	StatMemory.Add(int64(pages * pageSize))
 	return
 }
-- 
cgit mrf-deployment