@@ -4,6 +4,7 @@ package template
44
55import (
66 "testing"
7+ "time"
78
89 "github.com/google/uuid"
910 "github.com/stretchr/testify/assert"
@@ -80,3 +81,162 @@ func TestStorageTemplate_SchedulingMetadataSkipsPendingMemfile(t *testing.T) {
8081 assert .Empty (t , md .GetMemfileBaseBuildId ())
8182 assert .Equal (t , rootfsBase .String (), md .GetRootfsBaseBuildId ())
8283}
84+
85+ // footprintHeader builds a header whose mapping has exactly n entries, using a
86+ // distinct build id per entry so nothing merges.
87+ func footprintHeader (t * testing.T , n int ) * header.Header {
88+ t .Helper ()
89+
90+ const blockSize = uint64 (4096 )
91+
92+ maps := make ([]header.BuildMap , n )
93+ for i := range maps {
94+ maps [i ] = header.BuildMap {
95+ Offset : uint64 (i ) * blockSize ,
96+ Length : blockSize ,
97+ BuildId : uuid .New (),
98+ }
99+ }
100+
101+ h , err := header .NewHeader (& header.Metadata {
102+ Version : 3 ,
103+ BlockSize : blockSize ,
104+ Size : uint64 (n ) * blockSize ,
105+ BuildId : uuid .New (),
106+ BaseBuildId : uuid .New (),
107+ }, maps )
108+ require .NoError (t , err )
109+ require .Equal (t , n , h .Mapping .Len ())
110+
111+ return h
112+ }
113+
114+ // The gauge exists to size a bound, so it has to see every mapping the template
115+ // keeps alive. After a pause the provisional header stays referenced by
116+ // memfileHeader while the device has already moved on to the deduped one, so a
117+ // gauge that reads only the devices reports one mapping where two are resident
118+ // — understating exactly the footprint it is used to size.
119+ func TestStorageTemplate_HeaderFootprintCountsRetainedHolders (t * testing.T ) {
120+ t .Parallel ()
121+
122+ provisional := footprintHeader (t , 12 )
123+ deduped := footprintHeader (t , 3 )
124+ rootfsHdr := footprintHeader (t , 5 )
125+
126+ memDev := blockmocks .NewMockReadonlyDevice (t )
127+ memDev .EXPECT ().Header ().Return (deduped )
128+ rootfsDev := blockmocks .NewMockReadonlyDevice (t )
129+ rootfsDev .EXPECT ().Header ().Return (rootfsHdr )
130+
131+ tmpl := & storageTemplate {
132+ memfile : utils .NewSetOnce [block.ReadonlyDevice ](),
133+ rootfs : utils .NewSetOnce [block.ReadonlyDevice ](),
134+ memfileHeader : resolvedHeader (provisional ),
135+ rootfsHeader : resolvedHeader (rootfsHdr ),
136+ durableMemfileHeader : resolvedHeader (deduped ),
137+ }
138+ require .NoError (t , tmpl .memfile .SetValue (memDev ))
139+ require .NoError (t , tmpl .rootfs .SetValue (rootfsDev ))
140+
141+ entries , bytes := tmpl .headerFootprint ()
142+
143+ wantEntries := provisional .Mapping .Len () + deduped .Mapping .Len () + rootfsHdr .Mapping .Len ()
144+ wantBytes := provisional .Mapping .ByteSize () + deduped .Mapping .ByteSize () + rootfsHdr .Mapping .ByteSize ()
145+
146+ assert .Equal (t , wantEntries , entries , "the provisional header is still resident and must be counted" )
147+ assert .Equal (t , wantBytes , bytes )
148+ }
149+
150+ // The holders usually resolve to the same headers the devices carry. Counting
151+ // those twice would overstate the number the sizing decisions are made from.
152+ func TestStorageTemplate_HeaderFootprintCountsEachMappingOnce (t * testing.T ) {
153+ t .Parallel ()
154+
155+ memHdr := footprintHeader (t , 7 )
156+ rootfsHdr := footprintHeader (t , 4 )
157+
158+ memDev := blockmocks .NewMockReadonlyDevice (t )
159+ memDev .EXPECT ().Header ().Return (memHdr )
160+ rootfsDev := blockmocks .NewMockReadonlyDevice (t )
161+ rootfsDev .EXPECT ().Header ().Return (rootfsHdr )
162+
163+ tmpl := & storageTemplate {
164+ memfile : utils .NewSetOnce [block.ReadonlyDevice ](),
165+ rootfs : utils .NewSetOnce [block.ReadonlyDevice ](),
166+ memfileHeader : resolvedHeader (memHdr ),
167+ rootfsHeader : resolvedHeader (rootfsHdr ),
168+ }
169+ require .NoError (t , tmpl .memfile .SetValue (memDev ))
170+ require .NoError (t , tmpl .rootfs .SetValue (rootfsDev ))
171+
172+ entries , bytes := tmpl .headerFootprint ()
173+
174+ assert .Equal (t , memHdr .Mapping .Len ()+ rootfsHdr .Mapping .Len (), entries )
175+ assert .Equal (t , memHdr .Mapping .ByteSize ()+ rootfsHdr .Mapping .ByteSize (), bytes )
176+ }
177+
178+ // A template that is still fetching must contribute nothing rather than block
179+ // the metrics collection goroutine on its unresolved futures.
180+ func TestStorageTemplate_HeaderFootprintSkipsUnresolved (t * testing.T ) {
181+ t .Parallel ()
182+
183+ tmpl := & storageTemplate {
184+ memfile : utils .NewSetOnce [block.ReadonlyDevice ](),
185+ rootfs : utils .NewSetOnce [block.ReadonlyDevice ](),
186+ memfileHeader : utils .NewSetOnce [* header.Header ](),
187+ rootfsHeader : utils .NewSetOnce [* header.Header ](),
188+ }
189+
190+ done := make (chan struct {})
191+ go func () {
192+ defer close (done )
193+
194+ entries , bytes := tmpl .headerFootprint ()
195+ assert .Zero (t , entries )
196+ assert .Zero (t , bytes )
197+ }()
198+
199+ select {
200+ case <- done :
201+ case <- time .After (5 * time .Second ):
202+ t .Fatal ("headerFootprint blocked on an unresolved template" )
203+ }
204+ }
205+
206+ // One allocation reaches headerFootprint by two routes once a snapshot's upload
207+ // publishes: CloneForUpload copies the Header struct, so the clone shares the
208+ // source's Mapping slices, and publish installs the clone on the device while
209+ // the holder keeps the source. Header identity does not see that, so the gauge
210+ // would step up as uploads land — reading like retention growth rather than a
211+ // counting artifact, on exactly the population a byte budget gets sized from.
212+ func TestStorageTemplate_HeaderFootprintCountsSharedMappingOnce (t * testing.T ) {
213+ t .Parallel ()
214+
215+ source := footprintHeader (t , 9 )
216+ published := source .CloneForUpload (source .Metadata .Version + 1 )
217+ require .NotSame (t , source , published , "the clone is a distinct header" )
218+ require .True (t , source .Mapping .SharesStorageWith (published .Mapping ),
219+ "the clone must share the source's mapping, or this test proves nothing" )
220+
221+ rootfsHdr := footprintHeader (t , 4 )
222+
223+ memDev := blockmocks .NewMockReadonlyDevice (t )
224+ memDev .EXPECT ().Header ().Return (published )
225+ rootfsDev := blockmocks .NewMockReadonlyDevice (t )
226+ rootfsDev .EXPECT ().Header ().Return (rootfsHdr )
227+
228+ tmpl := & storageTemplate {
229+ memfile : utils .NewSetOnce [block.ReadonlyDevice ](),
230+ rootfs : utils .NewSetOnce [block.ReadonlyDevice ](),
231+ memfileHeader : resolvedHeader (source ),
232+ rootfsHeader : resolvedHeader (rootfsHdr ),
233+ }
234+ require .NoError (t , tmpl .memfile .SetValue (memDev ))
235+ require .NoError (t , tmpl .rootfs .SetValue (rootfsDev ))
236+
237+ entries , bytes := tmpl .headerFootprint ()
238+
239+ assert .Equal (t , source .Mapping .Len ()+ rootfsHdr .Mapping .Len (), entries ,
240+ "the shared mapping must be counted once, not once per header" )
241+ assert .Equal (t , source .Mapping .ByteSize ()+ rootfsHdr .Mapping .ByteSize (), bytes )
242+ }
0 commit comments