Provide better snapshot interface

Author: valerian-roche

pkg/cache/internal/cached_resource.go

@@ -0,0 +1,219 @@
+package internal
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"sync"
+	"time"
+
+	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
+
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+// Resource is the base interface for the xDS payload.
+type Resource interface {
+	proto.Message
+}
+
+// ResourceWithTTL is a Resource with an optional TTL.
+type ResourceWithTTL struct {
+	Resource Resource
+	TTL      *time.Duration
+}
+
+// CachedResource is used to track resources added by the user in the cache.
+// It contains the resource itself and its associated version (currently in two different modes).
+// It should not be altered once created, to allow concurrent access.
+type CachedResource struct {
+	Name    string
+	typeURL string
+
+	resource Resource
+	ttl      *time.Duration
+
+	// cacheVersion is the version of the cache at the time of last update, used in sotw.
+	cacheVersion string
+
+	marshalFunc                func() ([]byte, error)
+	computeResourceVersionFunc func() (string, error)
+}
+
+type CachedResourceOption = func(*CachedResource)
+
+// WithCacheVersion allows specifying the cacheVersion when the resource is set.
+func WithCacheVersion(version string) CachedResourceOption {
+	return func(r *CachedResource) { r.cacheVersion = version }
+}
+
+// WithMarshaledResource enables the user to provide the already marshaled bytes if they have them.
+// Those bytes should strive at being consistent if the object has not changed (beware protobuf non-deterministic marshaling)
+// or alternatively the resource version should also then be set.
+// By default it is computed by performing a deterministic protobuf marshaling.
+func WithMarshaledResource(bytes []byte) CachedResourceOption {
+	if len(bytes) == 0 {
+		return func(*CachedResource) {}
+	}
+	return func(r *CachedResource) { r.marshalFunc = func() ([]byte, error) { return bytes, nil } }
+}
+
+// WithResourceVersion enables the user to provide the resource version to be used.
+// This version should be constant if the object has not changed to avoid needlessly sending resources to clients.
+// By default it is computed by hashing the serialized version of the resource.
+func WithResourceVersion(version string) CachedResourceOption {
+	if version == "" {
+		return func(*CachedResource) {}
+	}
+	return func(r *CachedResource) { r.computeResourceVersionFunc = func() (string, error) { return version, nil } }
+}
+
+// WithResourceTTL sets a TTL on the resource, that will be sent to the client with the payload.
+func WithResourceTTL(ttl *time.Duration) CachedResourceOption {
+	return func(r *CachedResource) { r.ttl = ttl }
+}
+
+func NewCachedResource(name, typeURL string, res Resource, opts ...CachedResourceOption) *CachedResource {
+	cachedRes := &CachedResource{
+		Name:     name,
+		typeURL:  typeURL,
+		resource: res,
+	}
+	for _, opt := range opts {
+		opt(cachedRes)
+	}
+	if cachedRes.marshalFunc == nil {
+		cachedRes.marshalFunc = sync.OnceValues(func() ([]byte, error) {
+			return marshalResource(res)
+		})
+	}
+	if cachedRes.computeResourceVersionFunc == nil {
+		cachedRes.computeResourceVersionFunc = sync.OnceValues(func() (string, error) {
+			marshaled, err := cachedRes.marshalFunc()
+			if err != nil {
+				return "", fmt.Errorf("marshaling resource: %w", err)
+			}
+			return hashResource(marshaled), nil
+		})
+	}
+	return cachedRes
+}
+
+// SetCacheVersion updates the cache version. This violates the assumption that all fields can be safely read concurrently.
+// It's required today for the linear cache constructor, where we are guaranteed resources are not being used concurently,
+// and should not be used elsewhere.
+func (c *CachedResource) SetCacheVersion(version string) {
+	c.cacheVersion = version
+}
+
+// HasTTL returns whether the resource has a TTL set.
+func (c *CachedResource) HasTTL() bool {
+	return c.ttl != nil
+}
+
+// getMarshaledResource lazily marshals the resource and returns the bytes.
+func (c *CachedResource) getMarshaledResource() ([]byte, error) {
+	return c.marshalFunc()
+}
+
+// GetResourceVersion returns a stable version reflecting the resource content.
+// By default it is built by hashing the serialized version of the object, using deterministic serializing.
+func (c *CachedResource) GetResourceVersion() (string, error) {
+	return c.computeResourceVersionFunc()
+}
+
+// GetVersion returns the version for the resource.
+// By default it returns the cache version when the resource was added, but if requested it will return a
+// version specific to the resource content.
+func (c *CachedResource) GetVersion(useResourceVersion bool) (string, error) {
+	if !useResourceVersion {
+		return c.cacheVersion, nil
+	}
+
+	return c.GetResourceVersion()
+}
+
+// GetRawResource returns the underlying resource for use in legacy accessors.
+func (c *CachedResource) GetRawResource() ResourceWithTTL {
+	return ResourceWithTTL{
+		Resource: c.resource,
+		TTL:      c.ttl,
+	}
+}
+
+var deltaResourceTypeURL = "type.googleapis.com/" + string(proto.MessageName(&discovery.Resource{}))
+
+// getResourceVersion lazily hashes the resource and returns the stable hash used to track version changes.
+func (c *CachedResource) GetSotwResource(isHeartbeat bool) (*anypb.Any, error) {
+	buildResource := func() (*anypb.Any, error) {
+		marshaled, err := c.getMarshaledResource()
+		if err != nil {
+			return nil, fmt.Errorf("marshaling: %w", err)
+		}
+		return &anypb.Any{
+			TypeUrl: c.typeURL,
+			Value:   marshaled,
+		}, nil
+	}
+
+	if c.ttl == nil {
+		return buildResource()
+	}
+
+	wrappedResource := &discovery.Resource{
+		Name: c.Name,
+		Ttl:  durationpb.New(*c.ttl),
+	}
+
+	if !isHeartbeat {
+		rsrc, err := buildResource()
+		if err != nil {
+			return nil, err
+		}
+		wrappedResource.Resource = rsrc
+	}
+
+	marshaled, err := marshalResource(wrappedResource)
+	if err != nil {
+		return nil, fmt.Errorf("marshaling discovery resource: %w", err)
+	}
+
+	return &anypb.Any{
+		TypeUrl: deltaResourceTypeURL,
+		Value:   marshaled,
+	}, nil
+}
+
+// getResourceVersion lazily hashes the resource and returns the stable hash used to track version changes.
+func (c *CachedResource) GetDeltaResource() (*discovery.Resource, error) {
+	marshaled, err := c.getMarshaledResource()
+	if err != nil {
+		return nil, fmt.Errorf("marshaling: %w", err)
+	}
+	version, err := c.GetResourceVersion()
+	if err != nil {
+		return nil, fmt.Errorf("computing version: %w", err)
+	}
+	return &discovery.Resource{
+		Name: c.Name,
+		Resource: &anypb.Any{
+			TypeUrl: c.typeURL,
+			Value:   marshaled,
+		},
+		Version: version,
+	}, nil
+}
+
+// hashResource will take a resource and create a SHA256 hash sum out of the marshaled bytes
+func hashResource(resource []byte) string {
+	hasher := sha256.New()
+	hasher.Write(resource)
+	return hex.EncodeToString(hasher.Sum(nil))
+}
+
+// marshalResource converts the Resource to MarshaledResource.
+func marshalResource(resource Resource) ([]byte, error) {
+	return proto.MarshalOptions{Deterministic: true}.Marshal(resource)
+}

pkg/cache/internal/cached_resource_test.go

@@ -0,0 +1 @@
+package internal

pkg/cache/types/snapshot.go

@@ -0,0 +1,116 @@
+package types //nolint:revive // var-naming: avoid meaningless package names
+
+import (
+	"time"
+
+	"github.com/envoyproxy/go-control-plane/pkg/cache/internal"
+
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+// SnapshotResource represents a resource to be provided to caches.
+type SnapshotResource struct {
+	// Mandatory
+	Name string
+	// Mandatory
+	Resource Resource
+	// Optional
+	TTL *time.Duration
+
+	// Optional
+	Serialized *anypb.Any
+	// Optional
+	Version string
+}
+
+func (r *SnapshotResource) asCachedResource(typeURL, cacheVersion string) *internal.CachedResource {
+	var serialized []byte
+	if r.Serialized != nil {
+		serialized = r.Serialized.Value
+	}
+	return internal.NewCachedResource(r.Name, typeURL, r.Resource,
+		internal.WithCacheVersion(cacheVersion),
+		internal.WithResourceTTL(r.TTL),
+		internal.WithResourceVersion(r.Version),
+		internal.WithMarshaledResource(serialized))
+}
+
+// TypeSnapshot represents the resources for a given type, associated with an opaque version.
+// The snapshot (and associated resources) must not be modified once provided to the cache.
+// A given TypeSnapshot instance can be provided in multiple snapshots, including for multiple node ids, as long as it remains immutable.
+type TypeSnapshot struct {
+	typeURL   string
+	version   string
+	resources map[string]*internal.CachedResource
+}
+
+func NewTypeSnapshot(typeURL, version string, resources []SnapshotResource) TypeSnapshot {
+	s := TypeSnapshot{
+		typeURL:   typeURL,
+		version:   version,
+		resources: make(map[string]*internal.CachedResource, len(resources)),
+	}
+	for _, res := range resources {
+		s.resources[res.Name] = res.asCachedResource(typeURL, version)
+	}
+	return s
+}
+
+// GetVersion returns the version of the snapshot.
+// Multiple TypeSnapshots in a given Snapshot can have different versions.
+func (s TypeSnapshot) GetVersion() string {
+	return s.version
+}
+
+// GetResources returns the resources in the snapshot.
+// The map must not be modified by the caller.
+func (s TypeSnapshot) GetResources() map[string]*internal.CachedResource {
+	return s.resources
+}
+
+// Snapshot represents a consistent set of resources for multiple types.
+// Once provided to a cache it should not be altered in any way.
+type Snapshot struct {
+	// defaultVersion is the negative version returned if there is no snapshot set for the provided type.
+	defaultVersion string
+	resources      map[string]TypeSnapshot
+}
+
+// NewSnapshot creates a snapshot with a single version for all resource types.
+func NewSnapshot(version string, resources map[string][]SnapshotResource) (*Snapshot, error) {
+	s := &Snapshot{
+		defaultVersion: version,
+		resources:      make(map[string]TypeSnapshot, len(resources)),
+	}
+	for typeURL, res := range resources {
+		s.resources[typeURL] = NewTypeSnapshot(typeURL, version, res)
+	}
+	return s, nil
+}
+
+// NewSnapshotFromTypeSnapshots creates a snapshot from per-type snapshots.
+// TypeSnapshot instances can be shared across snapshots, but must not be altered in any way once provided to at least one snapshot.
+func NewSnapshotFromTypeSnapshots(version string, snapshots []TypeSnapshot) (*Snapshot, error) {
+	s := &Snapshot{
+		defaultVersion: version,
+		resources:      make(map[string]TypeSnapshot, len(snapshots)),
+	}
+	for _, snap := range snapshots {
+		s.resources[snap.typeURL] = snap
+	}
+	return s, nil
+}
+
+// GetVersion returns the current version of the resource indicated by typeURL.
+// The version string that is returned is opaque and should only be compared for equality.
+func (s *Snapshot) GetVersion(typeURL string) string {
+	typeSnapshot, ok := s.resources[typeURL]
+	if !ok {
+		return s.defaultVersion
+	}
+	return typeSnapshot.version
+}
+
+func (s *Snapshot) GetTypeSnapshot(typeURL string) TypeSnapshot {
+	return s.resources[typeURL]
+}

pkg/cache/types/types.go

@@ -1,31 +1,23 @@
 package types //nolint:revive // var-naming: avoid meaningless package names
 
 import (
-	"time"
+	"github.com/envoyproxy/go-control-plane/pkg/cache/internal"
 
 	"google.golang.org/protobuf/proto"
 )
 
 // Resource is the base interface for the xDS payload.
-type Resource interface {
-	proto.Message
-}
+type Resource = internal.Resource
 
 // ResourceWithTTL is a Resource with an optional TTL.
-type ResourceWithTTL struct {
-	Resource Resource
-	TTL      *time.Duration
-}
+type ResourceWithTTL = internal.ResourceWithTTL
 
 // ResourceWithName provides a name for out-of-tree resources.
 type ResourceWithName interface {
 	proto.Message
 	GetName() string
 }
 
-// MarshaledResource is an alias for the serialized binary array.
-type MarshaledResource = []byte
-
 // SkipFetchError is the error returned when the cache fetch is short
 // circuited due to the client's version already being up-to-date.
 type SkipFetchError struct{}