@@ -30,6 +30,7 @@ import (
3030 "github.com/containerd/containerd/v2/pkg/cri/util"
3131 "github.com/containerd/containerd/v2/platforms"
3232 docker "github.com/distribution/reference"
33+ "k8s.io/apimachinery/pkg/util/sets"
3334
3435 "github.com/opencontainers/go-digest"
3536 imagedigest "github.com/opencontainers/go-digest"
@@ -89,8 +90,9 @@ func NewStore(img images.Store, provider InfoProvider, platform platforms.MatchC
8990 provider : provider ,
9091 platform : platform ,
9192 store : & store {
92- images : make (map [string ]Image ),
93- digestSet : digestset .NewSet (),
93+ images : make (map [string ]Image ),
94+ digestSet : digestset .NewSet (),
95+ pinnedRefs : make (map [string ]sets.Set [string ]),
9496 },
9597 }
9698}
@@ -130,7 +132,13 @@ func (s *Store) update(ref string, img *Image) error {
130132 }
131133 if oldExist {
132134 if oldID == img .ID {
133- return nil
135+ if s .store .isPinned (img .ID , ref ) == img .Pinned {
136+ return nil
137+ }
138+ if img .Pinned {
139+ return s .store .pin (img .ID , ref )
140+ }
141+ return s .store .unpin (img .ID , ref )
134142 }
135143 // Updated. Remove tag from old image.
136144 s .store .delete (oldID , ref )
@@ -206,9 +214,10 @@ func (s *Store) List() []Image {
206214}
207215
208216type store struct {
209- lock sync.RWMutex
210- images map [string ]Image
211- digestSet * digestset.Set
217+ lock sync.RWMutex
218+ images map [string ]Image
219+ digestSet * digestset.Set
220+ pinnedRefs map [string ]sets.Set [string ]
212221}
213222
214223func (s * store ) list () []Image {
@@ -233,6 +242,14 @@ func (s *store) add(img Image) error {
233242 }
234243 }
235244
245+ if img .Pinned {
246+ if refs := s .pinnedRefs [img .ID ]; refs == nil {
247+ s .pinnedRefs [img .ID ] = sets .New (img .References ... )
248+ } else {
249+ refs .Insert (img .References ... )
250+ }
251+ }
252+
236253 i , ok := s .images [img .ID ]
237254 if ! ok {
238255 // If the image doesn't exist, add it.
@@ -246,6 +263,73 @@ func (s *store) add(img Image) error {
246263 return nil
247264}
248265
266+ func (s * store ) isPinned (id , ref string ) bool {
267+ s .lock .RLock ()
268+ defer s .lock .RUnlock ()
269+ digest , err := s .digestSet .Lookup (id )
270+ if err != nil {
271+ return false
272+ }
273+ refs := s .pinnedRefs [digest .String ()]
274+ return refs != nil && refs .Has (ref )
275+ }
276+
277+ func (s * store ) pin (id , ref string ) error {
278+ s .lock .Lock ()
279+ defer s .lock .Unlock ()
280+ digest , err := s .digestSet .Lookup (id )
281+ if err != nil {
282+ if err == digestset .ErrDigestNotFound {
283+ err = errdefs .ErrNotFound
284+ }
285+ return err
286+ }
287+ i , ok := s .images [digest .String ()]
288+ if ! ok {
289+ return errdefs .ErrNotFound
290+ }
291+
292+ if refs := s .pinnedRefs [digest .String ()]; refs == nil {
293+ s .pinnedRefs [digest .String ()] = sets .New (ref )
294+ } else {
295+ refs .Insert (ref )
296+ }
297+ i .Pinned = true
298+ s .images [digest .String ()] = i
299+ return nil
300+ }
301+
302+ func (s * store ) unpin (id , ref string ) error {
303+ s .lock .Lock ()
304+ defer s .lock .Unlock ()
305+ digest , err := s .digestSet .Lookup (id )
306+ if err != nil {
307+ if err == digestset .ErrDigestNotFound {
308+ err = errdefs .ErrNotFound
309+ }
310+ return err
311+ }
312+ i , ok := s .images [digest .String ()]
313+ if ! ok {
314+ return errdefs .ErrNotFound
315+ }
316+
317+ refs := s .pinnedRefs [digest .String ()]
318+ if refs == nil {
319+ return nil
320+ }
321+ if refs .Delete (ref ); len (refs ) > 0 {
322+ return nil
323+ }
324+
325+ // delete unpinned image, we only need to keep the pinned
326+ // entries in the map
327+ delete (s .pinnedRefs , digest .String ())
328+ i .Pinned = false
329+ s .images [digest .String ()] = i
330+ return nil
331+ }
332+
249333func (s * store ) get (id string ) (Image , error ) {
250334 s .lock .RLock ()
251335 defer s .lock .RUnlock ()
@@ -277,10 +361,20 @@ func (s *store) delete(id, ref string) {
277361 }
278362 i .References = util .SubtractStringSlice (i .References , ref )
279363 if len (i .References ) != 0 {
364+ if refs := s .pinnedRefs [digest .String ()]; refs != nil {
365+ if refs .Delete (ref ); len (refs ) == 0 {
366+ i .Pinned = false
367+ // delete unpinned image, we only need to keep the pinned
368+ // entries in the map
369+ delete (s .pinnedRefs , digest .String ())
370+ }
371+ }
372+
280373 s .images [digest .String ()] = i
281374 return
282375 }
283376 // Remove the image if it is not referenced any more.
284377 s .digestSet .Remove (digest )
285378 delete (s .images , digest .String ())
379+ delete (s .pinnedRefs , digest .String ())
286380}
0 commit comments