1package solver 2 3import ( 4 "context" 5 "sync" 6 "time" 7 8 "github.com/moby/buildkit/session" 9 "github.com/pkg/errors" 10) 11 12func NewInMemoryCacheStorage() CacheKeyStorage { 13 return &inMemoryStore{ 14 byID: map[string]*inMemoryKey{}, 15 byResult: map[string]map[string]struct{}{}, 16 } 17} 18 19type inMemoryStore struct { 20 mu sync.RWMutex 21 byID map[string]*inMemoryKey 22 byResult map[string]map[string]struct{} 23} 24 25type inMemoryKey struct { 26 id string 27 results map[string]CacheResult 28 links map[CacheInfoLink]map[string]struct{} 29 backlinks map[string]struct{} 30} 31 32func (s *inMemoryStore) Exists(id string) bool { 33 s.mu.RLock() 34 defer s.mu.RUnlock() 35 if k, ok := s.byID[id]; ok { 36 return len(k.links) > 0 || len(k.results) > 0 37 } 38 return false 39} 40 41func newInMemoryKey(id string) *inMemoryKey { 42 return &inMemoryKey{ 43 results: map[string]CacheResult{}, 44 links: map[CacheInfoLink]map[string]struct{}{}, 45 backlinks: map[string]struct{}{}, 46 id: id, 47 } 48} 49 50func (s *inMemoryStore) Walk(fn func(string) error) error { 51 s.mu.RLock() 52 ids := make([]string, 0, len(s.byID)) 53 for id := range s.byID { 54 ids = append(ids, id) 55 } 56 s.mu.RUnlock() 57 58 for _, id := range ids { 59 if err := fn(id); err != nil { 60 return err 61 } 62 } 63 return nil 64} 65 66func (s *inMemoryStore) WalkResults(id string, fn func(CacheResult) error) error { 67 s.mu.RLock() 68 69 k, ok := s.byID[id] 70 if !ok { 71 s.mu.RUnlock() 72 return nil 73 } 74 copy := make([]CacheResult, 0, len(k.results)) 75 for _, res := range k.results { 76 copy = append(copy, res) 77 } 78 s.mu.RUnlock() 79 80 for _, res := range copy { 81 if err := fn(res); err != nil { 82 return err 83 } 84 } 85 return nil 86} 87 88func (s *inMemoryStore) Load(id string, resultID string) (CacheResult, error) { 89 s.mu.RLock() 90 defer s.mu.RUnlock() 91 k, ok := s.byID[id] 92 if !ok { 93 return CacheResult{}, errors.Wrapf(ErrNotFound, "no such key %s", id) 94 } 95 r, ok := k.results[resultID] 96 if !ok { 97 return CacheResult{}, errors.WithStack(ErrNotFound) 98 } 99 return r, nil 100} 101 102func (s *inMemoryStore) AddResult(id string, res CacheResult) error { 103 s.mu.Lock() 104 defer s.mu.Unlock() 105 k, ok := s.byID[id] 106 if !ok { 107 k = newInMemoryKey(id) 108 s.byID[id] = k 109 } 110 k.results[res.ID] = res 111 m, ok := s.byResult[res.ID] 112 if !ok { 113 m = map[string]struct{}{} 114 s.byResult[res.ID] = m 115 } 116 m[id] = struct{}{} 117 return nil 118} 119 120func (s *inMemoryStore) WalkIDsByResult(resultID string, fn func(string) error) error { 121 s.mu.Lock() 122 123 ids := map[string]struct{}{} 124 for id := range s.byResult[resultID] { 125 ids[id] = struct{}{} 126 } 127 s.mu.Unlock() 128 129 for id := range ids { 130 if err := fn(id); err != nil { 131 return err 132 } 133 } 134 135 return nil 136} 137 138func (s *inMemoryStore) Release(resultID string) error { 139 s.mu.Lock() 140 defer s.mu.Unlock() 141 142 ids, ok := s.byResult[resultID] 143 if !ok { 144 return nil 145 } 146 147 for id := range ids { 148 k, ok := s.byID[id] 149 if !ok { 150 continue 151 } 152 153 delete(k.results, resultID) 154 delete(s.byResult[resultID], id) 155 if len(s.byResult[resultID]) == 0 { 156 delete(s.byResult, resultID) 157 } 158 159 s.emptyBranchWithParents(k) 160 } 161 162 return nil 163} 164 165func (s *inMemoryStore) emptyBranchWithParents(k *inMemoryKey) { 166 if len(k.results) != 0 || len(k.links) != 0 { 167 return 168 } 169 for id := range k.backlinks { 170 p, ok := s.byID[id] 171 if !ok { 172 continue 173 } 174 for l := range p.links { 175 delete(p.links[l], k.id) 176 if len(p.links[l]) == 0 { 177 delete(p.links, l) 178 } 179 } 180 s.emptyBranchWithParents(p) 181 } 182 183 delete(s.byID, k.id) 184} 185 186func (s *inMemoryStore) AddLink(id string, link CacheInfoLink, target string) error { 187 s.mu.Lock() 188 defer s.mu.Unlock() 189 k, ok := s.byID[id] 190 if !ok { 191 k = newInMemoryKey(id) 192 s.byID[id] = k 193 } 194 k2, ok := s.byID[target] 195 if !ok { 196 k2 = newInMemoryKey(target) 197 s.byID[target] = k2 198 } 199 m, ok := k.links[link] 200 if !ok { 201 m = map[string]struct{}{} 202 k.links[link] = m 203 } 204 205 k2.backlinks[id] = struct{}{} 206 m[target] = struct{}{} 207 return nil 208} 209 210func (s *inMemoryStore) WalkLinks(id string, link CacheInfoLink, fn func(id string) error) error { 211 s.mu.RLock() 212 k, ok := s.byID[id] 213 if !ok { 214 s.mu.RUnlock() 215 return nil 216 } 217 var links []string 218 for target := range k.links[link] { 219 links = append(links, target) 220 } 221 s.mu.RUnlock() 222 223 for _, t := range links { 224 if err := fn(t); err != nil { 225 return err 226 } 227 } 228 return nil 229} 230 231func (s *inMemoryStore) HasLink(id string, link CacheInfoLink, target string) bool { 232 s.mu.RLock() 233 defer s.mu.RUnlock() 234 if k, ok := s.byID[id]; ok { 235 if v, ok := k.links[link]; ok { 236 if _, ok := v[target]; ok { 237 return true 238 } 239 } 240 } 241 return false 242} 243 244func (s *inMemoryStore) WalkBacklinks(id string, fn func(id string, link CacheInfoLink) error) error { 245 s.mu.RLock() 246 k, ok := s.byID[id] 247 if !ok { 248 s.mu.RUnlock() 249 return nil 250 } 251 var outIDs []string 252 var outLinks []CacheInfoLink 253 for bid := range k.backlinks { 254 b, ok := s.byID[bid] 255 if !ok { 256 continue 257 } 258 for l, m := range b.links { 259 if _, ok := m[id]; !ok { 260 continue 261 } 262 outIDs = append(outIDs, bid) 263 outLinks = append(outLinks, CacheInfoLink{ 264 Digest: rootKey(l.Digest, l.Output), 265 Input: l.Input, 266 Selector: l.Selector, 267 }) 268 } 269 } 270 s.mu.RUnlock() 271 272 for i := range outIDs { 273 if err := fn(outIDs[i], outLinks[i]); err != nil { 274 return err 275 } 276 } 277 return nil 278} 279 280func NewInMemoryResultStorage() CacheResultStorage { 281 return &inMemoryResultStore{m: &sync.Map{}} 282} 283 284type inMemoryResultStore struct { 285 m *sync.Map 286} 287 288func (s *inMemoryResultStore) Save(r Result, createdAt time.Time) (CacheResult, error) { 289 s.m.Store(r.ID(), r) 290 return CacheResult{ID: r.ID(), CreatedAt: createdAt}, nil 291} 292 293func (s *inMemoryResultStore) Load(ctx context.Context, res CacheResult) (Result, error) { 294 v, ok := s.m.Load(res.ID) 295 if !ok { 296 return nil, errors.WithStack(ErrNotFound) 297 } 298 return v.(Result), nil 299} 300 301func (s *inMemoryResultStore) LoadRemote(_ context.Context, _ CacheResult, _ session.Group) (*Remote, error) { 302 return nil, nil 303} 304 305func (s *inMemoryResultStore) Exists(id string) bool { 306 _, ok := s.m.Load(id) 307 return ok 308} 309