Skip to content

Commit

Permalink
Merge pull request #33 from stacklok/cache-references
Browse files Browse the repository at this point in the history
Add cache implementation
  • Loading branch information
JAORMX committed Dec 11, 2023
2 parents 41c54fa + 85ed3a5 commit be7bc4d
Show file tree
Hide file tree
Showing 7 changed files with 123 additions and 10 deletions.
3 changes: 2 additions & 1 deletion cmd/containerimage/yamlreplacer.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ func (r *yamlReplacer) do(ctx context.Context, _ *config.Config) error {
modified.Store(false)

var eg errgroup.Group
cache := utils.NewRefCacher()

err := utils.Traverse(bfs, base, func(path string, info fs.FileInfo) error {
eg.Go(func() error {
Expand All @@ -68,7 +69,7 @@ func (r *yamlReplacer) do(ctx context.Context, _ *config.Config) error {
r.Logf("Processing %s\n", path)

buf := bytes.Buffer{}
m, err := containers.ReplaceReferenceFromYAML(ctx, r.imageRegex, f, &buf)
m, err := containers.ReplaceReferenceFromYAMLWithCache(ctx, r.imageRegex, f, &buf, cache)
if err != nil {
return fmt.Errorf("failed to process YAML file %s: %w", path, err)
}
Expand Down
3 changes: 2 additions & 1 deletion cmd/ghactions/replacer.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,12 @@ func (r *replacer) do(ctx context.Context, cfg *config.Config) error {

// error group
var eg errgroup.Group
cache := utils.NewRefCacher()

err := ghactions.TraverseGitHubActionWorkflows(bfs, base, func(path string, wflow *yaml.Node) error {
eg.Go(func() error {
r.Logf("Processing %s\n", path)
m, err := ghactions.ModifyReferencesInYAML(ctx, r.restIf, wflow, &cfg.GHActions)
m, err := ghactions.ModifyReferencesInYAMLWithCache(ctx, r.restIf, wflow, &cfg.GHActions, cache)
if err != nil {
return fmt.Errorf("failed to process YAML file %s: %w", path, err)
}
Expand Down
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ require (
github.com/go-git/go-billy/v5 v5.5.0
github.com/google/go-containerregistry v0.17.0
github.com/google/go-github/v56 v56.0.0
github.com/puzpuzpuz/xsync v1.5.2
github.com/spf13/cobra v1.8.0
github.com/stretchr/testify v1.8.4
golang.org/x/sync v0.2.0
gopkg.in/yaml.v3 v3.0.1
)

Expand All @@ -31,6 +33,5 @@ require (
github.com/sirupsen/logrus v1.9.1 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/vbatts/tar-split v0.11.3 // indirect
golang.org/x/sync v0.2.0 // indirect
golang.org/x/sys v0.12.0 // indirect
)
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/puzpuzpuz/xsync v1.5.2 h1:yRAP4wqSOZG+/4pxJ08fPTwrfL0IzE/LKQ/cw509qGY=
github.com/puzpuzpuz/xsync v1.5.2/go.mod h1:K98BYhX3k1dQ2M63t1YNVDanbwUPmBCAhNmVrrxfiGg=
github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
Expand Down
23 changes: 20 additions & 3 deletions pkg/containers/replace.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ import (
"regexp"

"github.com/google/go-containerregistry/pkg/name"

"github.com/stacklok/frizbee/pkg/utils"
)

// ReplaceImageReferenceFromYAML replaces the image reference in the input text with the digest
Expand All @@ -32,6 +34,14 @@ func ReplaceImageReferenceFromYAML(ctx context.Context, input io.Reader, output

// ReplaceReferenceFromYAML replaces the image reference in the input text with the digest
func ReplaceReferenceFromYAML(ctx context.Context, keyRegex string, input io.Reader, output io.Writer) (bool, error) {
cache := utils.NewUnsafeCacher()
return ReplaceReferenceFromYAMLWithCache(ctx, keyRegex, input, output, cache)
}

// ReplaceReferenceFromYAMLWithCache replaces the image reference in the input text with the digest
// and uses the provided cache to store the digests.
func ReplaceReferenceFromYAMLWithCache(
ctx context.Context, keyRegex string, input io.Reader, output io.Writer, cache utils.RefCacher) (bool, error) {
scanner := bufio.NewScanner(input)
re, err := regexp.Compile(fmt.Sprintf(`(\s*%s):\s*([^\s]+)`, keyRegex))
if err != nil {
Expand All @@ -54,9 +64,16 @@ func ReplaceReferenceFromYAML(ctx context.Context, keyRegex string, input io.Rea
return match
}

digest, err := GetDigestFromRef(ctx, ref)
if err != nil {
return match
var digest string
if d, ok := cache.Load(ref.Identifier()); ok {
digest = d
} else {
digest, err = GetDigestFromRef(ctx, ref)
if err != nil {
return match
}

cache.Store(ref.Identifier(), digest)
}

imgWithoutTag := ref.Context().Name()
Expand Down
29 changes: 25 additions & 4 deletions pkg/ghactions/ghactions.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import (

"github.com/stacklok/frizbee/pkg/config"
"github.com/stacklok/frizbee/pkg/interfaces"
"github.com/stacklok/frizbee/pkg/utils"
)

// IsLocal returns true if the input is a local path.
Expand Down Expand Up @@ -81,6 +82,17 @@ func GetChecksum(ctx context.Context, restIf interfaces.REST, action, ref string
// Note that the given YAML structure is modified in-place.
// The function returns true if any references were modified.
func ModifyReferencesInYAML(ctx context.Context, restIf interfaces.REST, node *yaml.Node, cfg *config.GHActions) (bool, error) {
cache := utils.NewUnsafeCacher()
return ModifyReferencesInYAMLWithCache(ctx, restIf, node, cfg, cache)
}

// ModifyReferencesInYAMLWithCache takes the given YAML structure and replaces
// all references to tags with the checksum of the tag.
// Note that the given YAML structure is modified in-place.
// The function returns true if any references were modified.
// The function uses the provided cache to store the checksums.
func ModifyReferencesInYAMLWithCache(
ctx context.Context, restIf interfaces.REST, node *yaml.Node, cfg *config.GHActions, cache utils.RefCacher) (bool, error) {
// `uses` will be immediately before the action
// name in the YAML `Content` array. We use a toggle
// to track if we've found `uses` and then look for
Expand Down Expand Up @@ -111,9 +123,18 @@ func ModifyReferencesInYAML(ctx context.Context, restIf interfaces.REST, node *y
return modified, fmt.Errorf("failed to parse action reference '%s': %w", v.Value, err)
}

sum, err := GetChecksum(ctx, restIf, act, ref)
if err != nil {
return modified, fmt.Errorf("failed to get checksum for action '%s': %w", v.Value, err)
var sum string

// Check if we have a cached value
if val, ok := cache.Load(v.Value); ok {
sum = val
} else {
sum, err = GetChecksum(ctx, restIf, act, ref)
if err != nil {
return modified, fmt.Errorf("failed to get checksum for action '%s': %w", v.Value, err)
}

cache.Store(v.Value, sum)
}

if ref != sum {
Expand All @@ -125,7 +146,7 @@ func ModifyReferencesInYAML(ctx context.Context, restIf interfaces.REST, node *y
}

// Otherwise recursively look more
m, err := ModifyReferencesInYAML(ctx, restIf, v, cfg)
m, err := ModifyReferencesInYAMLWithCache(ctx, restIf, v, cfg, cache)
if err != nil {
return m, err
}
Expand Down
70 changes: 70 additions & 0 deletions pkg/utils/cache.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
//
// Copyright 2023 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package utils

import (
"github.com/puzpuzpuz/xsync"
)

// RefCacher is an interface for caching references.
type RefCacher interface {
Store(key, value string)
Load(key string) (string, bool)
}

type refCacher struct {
cache *xsync.MapOf[string, string]
}

// NewRefCacher returns a new RefCacher. The default implementation is
// thread-safe.
func NewRefCacher() RefCacher {
return &refCacher{
cache: xsync.NewMapOf[string](),
}
}

// Store stores a key-value pair.
func (r *refCacher) Store(key, value string) {
r.cache.Store(key, value)
}

// Load loads a value for a given key.
func (r *refCacher) Load(key string) (string, bool) {
return r.cache.Load(key)
}

type unsafeCacher struct {
cache map[string]string
}

// NewUnsafeCacher returns a new RefCacher that's not thread-safe.
func NewUnsafeCacher() RefCacher {
return &unsafeCacher{
cache: map[string]string{},
}
}

// Store stores a key-value pair.
func (r *unsafeCacher) Store(key, value string) {
r.cache[key] = value
}

// Load loads a value for a given key.
func (r *unsafeCacher) Load(key string) (string, bool) {
v, ok := r.cache[key]
return v, ok
}

0 comments on commit be7bc4d

Please sign in to comment.