buildkit/solver-next/cache.go

566 lines
13 KiB
Go

package solver
import (
"context"
"fmt"
"sort"
"strings"
"sync"
"github.com/moby/buildkit/identity"
digest "github.com/opencontainers/go-digest"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
)
type internalMemoryKeyT string
var internalMemoryKey = internalMemoryKeyT("buildkit/memory-cache-id")
var NoSelector = digest.FromBytes(nil)
func NewInMemoryCacheManager() CacheManager {
return NewCacheManager(identity.NewID(), NewInMemoryCacheStorage(), NewInMemoryResultStorage())
}
func NewCacheManager(id string, storage CacheKeyStorage, results CacheResultStorage) CacheManager {
cm := &inMemoryCacheManager{
id: id,
backend: storage,
results: results,
}
storage.Walk(func(id string) error {
return storage.WalkResults(id, func(cr CacheResult) error {
if !results.Exists(cr.ID) {
storage.Release(cr.ID)
}
return nil
})
})
return cm
}
type inMemoryCacheKey struct {
manager *inMemoryCacheManager
cacheResult CacheResult
deps []CacheKeyWithSelector // only []*inMemoryCacheKey
digest digest.Digest
output Index
id string
CacheKey
}
func (ck *inMemoryCacheKey) Deps() []CacheKeyWithSelector {
return ck.deps
}
func (ck *inMemoryCacheKey) Digest() digest.Digest {
return ck.digest
}
func (ck *inMemoryCacheKey) Output() Index {
return ck.output
}
func withExporter(ck *inMemoryCacheKey, cacheResult *CacheResult, deps []CacheKeyWithSelector) ExportableCacheKey {
return ExportableCacheKey{ck, &cacheExporter{
inMemoryCacheKey: ck,
cacheResult: cacheResult,
deps: deps,
}}
}
type cacheExporter struct {
*inMemoryCacheKey
cacheResult *CacheResult
deps []CacheKeyWithSelector
}
func (ce *cacheExporter) Export(ctx context.Context, m map[digest.Digest]*ExportRecord, converter func(context.Context, Result) (*Remote, error)) (*ExportRecord, error) {
var res Result
if ce.cacheResult == nil {
cr, err := ce.inMemoryCacheKey.manager.getBestResult(ce.inMemoryCacheKey.id)
if err != nil {
return nil, err
}
ce.cacheResult = cr
}
var remote *Remote
var err error
if ce.cacheResult != nil {
remote, err = ce.inMemoryCacheKey.manager.results.LoadRemote(ctx, *ce.cacheResult)
if err != nil {
return nil, err
}
if remote == nil {
res, err = ce.inMemoryCacheKey.manager.results.Load(ctx, *ce.cacheResult)
if err != nil {
return nil, err
}
}
}
if res != nil && remote == nil {
remote, err = converter(ctx, res)
if err != nil {
return nil, err
}
}
cacheID := digest.FromBytes([]byte(ce.inMemoryCacheKey.id))
if remote != nil && len(remote.Descriptors) > 0 && remote.Descriptors[0].Digest != "" {
cacheID = remote.Descriptors[0].Digest
}
rec, ok := m[cacheID]
if !ok {
rec = &ExportRecord{
Digest: cacheID,
Remote: remote,
Links: make(map[CacheLink]struct{}),
}
m[cacheID] = rec
}
if len(ce.Deps()) == 0 {
rec.Links[CacheLink{
Output: ce.Output(),
Base: ce.Digest(),
}] = struct{}{}
}
for i, dep := range ce.deps {
if dep.CacheKey.Exporter == nil {
continue
}
r, err := dep.CacheKey.Export(ctx, m, converter)
if err != nil {
return nil, err
}
link := CacheLink{
Source: r.Digest,
Input: Index(i),
Output: ce.Output(),
Base: ce.Digest(),
Selector: dep.Selector,
}
rec.Links[link] = struct{}{}
}
return rec, nil
}
type inMemoryCacheManager struct {
mu sync.RWMutex
id string
backend CacheKeyStorage
results CacheResultStorage
}
func (c *inMemoryCacheManager) ID() string {
return c.id
}
func (c *inMemoryCacheManager) toInMemoryCacheKey(id string, dgst digest.Digest, output Index, deps []CacheKeyWithSelector) *inMemoryCacheKey {
ck := &inMemoryCacheKey{
id: id,
output: output,
digest: dgst,
manager: c,
CacheKey: NewCacheKey("", 0, nil),
deps: deps,
}
ck.SetValue(internalMemoryKey, id)
return ck
}
func (c *inMemoryCacheManager) getBestResult(id string) (*CacheResult, error) {
var results []*CacheResult
if err := c.backend.WalkResults(id, func(res CacheResult) error {
results = append(results, &res)
return nil
}); err != nil {
return nil, err
}
sort.Slice(results, func(i, j int) bool {
return results[i].CreatedAt.After(results[j].CreatedAt)
})
if len(results) > 0 {
return results[0], nil
}
return nil, nil
}
func (c *inMemoryCacheManager) Query(deps []CacheKeyWithSelector, input Index, dgst digest.Digest, output Index) ([]*CacheRecord, error) {
c.mu.RLock()
defer c.mu.RUnlock()
type dep struct {
results map[string]struct{}
key CacheKeyWithSelector
internalKey *inMemoryCacheKey
}
formatDeps := func(deps []dep, index int) []CacheKeyWithSelector {
keys := make([]CacheKeyWithSelector, index+1)
if len(deps) == 1 {
keys[index] = deps[0].key
} else {
k2 := make([]CacheKeyWithSelector, 0, len(deps))
for _, d := range deps {
k2 = append(k2, d.key)
}
keys[index] = CacheKeyWithSelector{CacheKey: ExportableCacheKey{CacheKey: NewCacheKey("", 0, k2)}}
}
return keys
}
allDeps := make([]dep, 0, len(deps))
for _, d := range deps {
for _, k := range c.getAllKeys(d) {
d := dep{key: k, results: map[string]struct{}{}}
internalKey, err := c.getInternalKey(k.CacheKey, false)
if err != nil {
if errors.Cause(err) == ErrNotFound {
allDeps = append(allDeps, d)
} else {
return nil, err
}
} else {
d.internalKey = internalKey
}
allDeps = append(allDeps, d)
}
}
allRes := map[string]struct{}{}
for _, d := range allDeps {
if d.internalKey != nil {
if err := c.backend.WalkLinks(d.internalKey.id, CacheInfoLink{input, output, dgst, d.key.Selector}, func(id string) error {
d.results[id] = struct{}{}
allRes[id] = struct{}{}
return nil
}); err != nil {
return nil, err
}
}
}
if len(deps) == 0 {
allRes[digest.FromBytes([]byte(fmt.Sprintf("%s@%d", dgst, output))).String()] = struct{}{}
}
outs := make([]*CacheRecord, 0, len(allRes))
for res := range allRes {
for _, d := range allDeps {
if d.internalKey == nil {
internalKey, err := c.getInternalKey(d.key.CacheKey, true)
if err != nil {
return nil, err
}
d.internalKey = internalKey
}
if _, ok := d.results[res]; !ok {
if err := c.backend.AddLink(d.internalKey.id, CacheInfoLink{
Input: input,
Output: output,
Digest: dgst,
Selector: d.key.Selector,
}, res); err != nil {
return nil, err
}
}
}
hadResults := false
fdeps := formatDeps(allDeps, int(input))
k := c.toInMemoryCacheKey(res, dgst, output, fdeps)
// TODO: invoke this only once per input
if err := c.backend.WalkResults(res, func(r CacheResult) error {
if c.results.Exists(r.ID) {
outs = append(outs, &CacheRecord{
ID: res + "@" + r.ID,
CacheKey: withExporter(k, &r, fdeps),
CacheManager: c,
Loadable: true,
CreatedAt: r.CreatedAt,
})
hadResults = true
} else {
c.backend.Release(r.ID)
}
return nil
}); err != nil {
return nil, err
}
if !hadResults {
if len(deps) == 0 {
if !c.backend.Exists(res) {
continue
}
}
outs = append(outs, &CacheRecord{
ID: res,
CacheKey: withExporter(k, nil, fdeps),
CacheManager: c,
Loadable: false,
})
}
}
return outs, nil
}
func (c *inMemoryCacheManager) Load(ctx context.Context, rec *CacheRecord) (Result, error) {
c.mu.RLock()
defer c.mu.RUnlock()
keyParts := strings.Split(rec.ID, "@")
if len(keyParts) != 2 {
return nil, errors.Errorf("invalid cache record ID")
}
ck, err := c.getInternalKey(rec.CacheKey, false)
if err != nil {
return nil, err
}
res, err := c.backend.Load(ck.id, keyParts[1])
if err != nil {
return nil, err
}
return c.results.Load(ctx, res)
}
func (c *inMemoryCacheManager) Save(k CacheKey, r Result) (ExportableCacheKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
empty := ExportableCacheKey{}
ck, err := c.getInternalKey(k, true)
if err != nil {
return empty, err
}
res, err := c.results.Save(r)
if err != nil {
return empty, err
}
if err := c.backend.AddResult(ck.id, res); err != nil {
return empty, err
}
return withExporter(ck, &res, ck.Deps()), nil
}
func (c *inMemoryCacheManager) getInternalKeys(d CacheKeyWithSelector, createIfNotExist bool) ([]CacheKeyWithSelector, error) {
keys := make([]CacheKeyWithSelector, 0, 1)
if d.CacheKey.Digest() == "" {
for _, d := range d.CacheKey.Deps() {
k, err := c.getInternalKey(d.CacheKey, createIfNotExist)
if err != nil {
if !createIfNotExist && errors.Cause(err) == ErrNotFound {
continue
}
return nil, err
}
keys = append(keys, CacheKeyWithSelector{Selector: d.Selector, CacheKey: ExportableCacheKey{CacheKey: k, Exporter: d.CacheKey.Exporter}})
}
} else {
k, err := c.getInternalKey(d.CacheKey, createIfNotExist)
if err != nil {
return nil, err
}
keys = append(keys, CacheKeyWithSelector{Selector: d.Selector, CacheKey: ExportableCacheKey{CacheKey: k, Exporter: d.CacheKey.Exporter}})
}
return keys, nil
}
func (c *inMemoryCacheManager) getAllKeys(d CacheKeyWithSelector) []CacheKeyWithSelector {
keys := make([]CacheKeyWithSelector, 0, 1)
if d.CacheKey.Digest() == "" {
for _, d := range d.CacheKey.Deps() {
keys = append(keys, d)
}
} else {
keys = append(keys, d)
}
return keys
}
func (c *inMemoryCacheManager) getInternalKey(k CacheKey, createIfNotExist bool) (*inMemoryCacheKey, error) {
if ck, ok := k.(ExportableCacheKey); ok {
k = ck.CacheKey
}
if ck, ok := k.(*inMemoryCacheKey); ok {
return ck, nil
}
internalV := k.GetValue(internalMemoryKey)
if internalV != nil {
return c.toInMemoryCacheKey(internalV.(string), k.Digest(), k.Output(), k.Deps()), nil
}
matches := make(map[string]struct{})
deps := make([][]CacheKeyWithSelector, 0, len(k.Deps()))
for i, inp := range k.Deps() {
allKeys := c.getAllKeys(inp)
cks := make([]CacheKeyWithSelector, 0, len(allKeys))
for _, k := range allKeys {
internalKey, err := c.getInternalKey(k.CacheKey, createIfNotExist)
if err == nil {
cks = append(cks, CacheKeyWithSelector{Selector: k.Selector, CacheKey: ExportableCacheKey{CacheKey: internalKey, Exporter: k.CacheKey.Exporter}})
}
}
if len(cks) == 0 {
return nil, errors.WithStack(ErrNotFound)
}
if i == 0 || len(matches) > 0 {
for _, ck := range cks {
internalCk := ck.CacheKey.CacheKey.(*inMemoryCacheKey)
m2 := make(map[string]struct{})
if err := c.backend.WalkLinks(internalCk.id, CacheInfoLink{
Input: Index(i),
Output: Index(k.Output()),
Digest: k.Digest(),
Selector: ck.Selector,
}, func(id string) error {
if i == 0 {
matches[id] = struct{}{}
} else {
m2[id] = struct{}{}
}
return nil
}); err != nil {
return nil, err
}
if i != 0 {
for id := range matches {
if _, ok := m2[id]; !ok {
delete(matches, id)
}
}
}
}
}
deps = append(deps, cks)
}
var internalKey string
if len(matches) == 0 && len(k.Deps()) > 0 {
if createIfNotExist {
internalKey = identity.NewID()
} else {
return nil, errors.WithStack(ErrNotFound)
}
} else {
for k := range matches {
internalKey = k
break
}
if len(k.Deps()) == 0 {
internalKey = digest.FromBytes([]byte(fmt.Sprintf("%s@%d", k.Digest(), k.Output()))).String()
}
return c.toInMemoryCacheKey(internalKey, k.Digest(), k.Output(), k.Deps()), nil
}
for i, dep := range deps {
for _, ck := range dep {
internalCk := ck.CacheKey.CacheKey.(*inMemoryCacheKey)
err := c.backend.AddLink(internalCk.id, CacheInfoLink{
Input: Index(i),
Output: k.Output(),
Digest: k.Digest(),
Selector: ck.Selector,
}, internalKey)
if err != nil {
return nil, err
}
}
}
return c.toInMemoryCacheKey(internalKey, k.Digest(), k.Output(), k.Deps()), nil
}
func newCombinedCacheManager(cms []CacheManager, main CacheManager) CacheManager {
return &combinedCacheManager{cms: cms, main: main}
}
type combinedCacheManager struct {
cms []CacheManager
main CacheManager
id string
idOnce sync.Once
}
func (cm *combinedCacheManager) ID() string {
cm.idOnce.Do(func() {
ids := make([]string, len(cm.cms))
for i, c := range cm.cms {
ids[i] = c.ID()
}
cm.id = digest.FromBytes([]byte(strings.Join(ids, ","))).String()
})
return cm.id
}
func (cm *combinedCacheManager) Query(inp []CacheKeyWithSelector, inputIndex Index, dgst digest.Digest, outputIndex Index) ([]*CacheRecord, error) {
eg, _ := errgroup.WithContext(context.TODO())
res := make(map[string]*CacheRecord, len(cm.cms))
var mu sync.Mutex
for i, c := range cm.cms {
func(i int, c CacheManager) {
eg.Go(func() error {
recs, err := c.Query(inp, inputIndex, dgst, outputIndex)
if err != nil {
return err
}
mu.Lock()
for _, r := range recs {
if _, ok := res[r.ID]; !ok || c == cm.main {
r.CacheManager = c
if c == cm.main {
r.Priority = 1
}
res[r.ID] = r
}
}
mu.Unlock()
return nil
})
}(i, c)
}
if err := eg.Wait(); err != nil {
return nil, err
}
out := make([]*CacheRecord, 0, len(res))
for _, r := range res {
out = append(out, r)
}
return out, nil
}
func (cm *combinedCacheManager) Load(ctx context.Context, rec *CacheRecord) (Result, error) {
return rec.CacheManager.Load(ctx, rec)
}
func (cm *combinedCacheManager) Save(key CacheKey, s Result) (ExportableCacheKey, error) {
return cm.main.Save(key, s)
}