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