@@ -27,6 +27,7 @@ import (
2727 "github.com/containerd/containerd/errdefs"
2828 "github.com/containerd/containerd/pkg/cri/labels"
2929 "github.com/containerd/containerd/pkg/cri/util"
30+ "k8s.io/apimachinery/pkg/util/sets"
3031
3132 imagedigest "github.com/opencontainers/go-digest"
3233 "github.com/opencontainers/go-digest/digestset"
@@ -68,8 +69,9 @@ func NewStore(client *containerd.Client) *Store {
6869 refCache : make (map [string ]string ),
6970 client : client ,
7071 store : & store {
71- images : make (map [string ]Image ),
72- digestSet : digestset .NewSet (),
72+ images : make (map [string ]Image ),
73+ digestSet : digestset .NewSet (),
74+ pinnedRefs : make (map [string ]sets.String ),
7375 },
7476 }
7577}
@@ -107,7 +109,13 @@ func (s *Store) update(ref string, img *Image) error {
107109 }
108110 if oldExist {
109111 if oldID == img .ID {
110- return nil
112+ if s .store .isPinned (img .ID , ref ) == img .Pinned {
113+ return nil
114+ }
115+ if img .Pinned {
116+ return s .store .pin (img .ID , ref )
117+ }
118+ return s .store .unpin (img .ID , ref )
111119 }
112120 // Updated. Remove tag from old image.
113121 s .store .delete (oldID , ref )
@@ -183,9 +191,10 @@ func (s *Store) List() []Image {
183191}
184192
185193type store struct {
186- lock sync.RWMutex
187- images map [string ]Image
188- digestSet * digestset.Set
194+ lock sync.RWMutex
195+ images map [string ]Image
196+ digestSet * digestset.Set
197+ pinnedRefs map [string ]sets.String
189198}
190199
191200func (s * store ) list () []Image {
@@ -210,6 +219,14 @@ func (s *store) add(img Image) error {
210219 }
211220 }
212221
222+ if img .Pinned {
223+ if refs := s .pinnedRefs [img .ID ]; refs == nil {
224+ s .pinnedRefs [img .ID ] = sets .NewString (img .References ... )
225+ } else {
226+ refs .Insert (img .References ... )
227+ }
228+ }
229+
213230 i , ok := s .images [img .ID ]
214231 if ! ok {
215232 // If the image doesn't exist, add it.
@@ -223,6 +240,73 @@ func (s *store) add(img Image) error {
223240 return nil
224241}
225242
243+ func (s * store ) isPinned (id , ref string ) bool {
244+ s .lock .RLock ()
245+ defer s .lock .RUnlock ()
246+ digest , err := s .digestSet .Lookup (id )
247+ if err != nil {
248+ return false
249+ }
250+ refs := s .pinnedRefs [digest .String ()]
251+ return refs != nil && refs .Has (ref )
252+ }
253+
254+ func (s * store ) pin (id , ref string ) error {
255+ s .lock .Lock ()
256+ defer s .lock .Unlock ()
257+ digest , err := s .digestSet .Lookup (id )
258+ if err != nil {
259+ if err == digestset .ErrDigestNotFound {
260+ err = errdefs .ErrNotFound
261+ }
262+ return err
263+ }
264+ i , ok := s .images [digest .String ()]
265+ if ! ok {
266+ return errdefs .ErrNotFound
267+ }
268+
269+ if refs := s .pinnedRefs [digest .String ()]; refs == nil {
270+ s .pinnedRefs [digest .String ()] = sets .NewString (ref )
271+ } else {
272+ refs .Insert (ref )
273+ }
274+ i .Pinned = true
275+ s .images [digest .String ()] = i
276+ return nil
277+ }
278+
279+ func (s * store ) unpin (id , ref string ) error {
280+ s .lock .Lock ()
281+ defer s .lock .Unlock ()
282+ digest , err := s .digestSet .Lookup (id )
283+ if err != nil {
284+ if err == digestset .ErrDigestNotFound {
285+ err = errdefs .ErrNotFound
286+ }
287+ return err
288+ }
289+ i , ok := s .images [digest .String ()]
290+ if ! ok {
291+ return errdefs .ErrNotFound
292+ }
293+
294+ refs := s .pinnedRefs [digest .String ()]
295+ if refs == nil {
296+ return nil
297+ }
298+ if refs .Delete (ref ); len (refs ) > 0 {
299+ return nil
300+ }
301+
302+ // delete unpinned image, we only need to keep the pinned
303+ // entries in the map
304+ delete (s .pinnedRefs , digest .String ())
305+ i .Pinned = false
306+ s .images [digest .String ()] = i
307+ return nil
308+ }
309+
226310func (s * store ) get (id string ) (Image , error ) {
227311 s .lock .RLock ()
228312 defer s .lock .RUnlock ()
@@ -254,10 +338,20 @@ func (s *store) delete(id, ref string) {
254338 }
255339 i .References = util .SubtractStringSlice (i .References , ref )
256340 if len (i .References ) != 0 {
341+ if refs := s .pinnedRefs [digest .String ()]; refs != nil {
342+ if refs .Delete (ref ); len (refs ) == 0 {
343+ i .Pinned = false
344+ // delete unpinned image, we only need to keep the pinned
345+ // entries in the map
346+ delete (s .pinnedRefs , digest .String ())
347+ }
348+ }
349+
257350 s .images [digest .String ()] = i
258351 return
259352 }
260353 // Remove the image if it is not referenced any more.
261354 s .digestSet .Remove (digest )
262355 delete (s .images , digest .String ())
356+ delete (s .pinnedRefs , digest .String ())
263357}
0 commit comments