diff --git a/.golangci.yml b/.golangci.yml index 656e94021f0..7d2d0d65f45 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -43,6 +43,10 @@ issues: - errcheck - goimports + - path: '^go/vt/vtadmin/cache/' + linters: + - structcheck + ### BEGIN: errcheck exclusion rules. Each rule should be considered # a TODO for removal after adding error checks to that package/file/etc, # except where otherwise noted. diff --git a/go/cmd/vtadmin/main.go b/go/cmd/vtadmin/main.go index 6e2bda475d9..0a908c47d5f 100644 --- a/go/cmd/vtadmin/main.go +++ b/go/cmd/vtadmin/main.go @@ -28,6 +28,7 @@ import ( "vitess.io/vitess/go/trace" "vitess.io/vitess/go/vt/log" "vitess.io/vitess/go/vt/vtadmin" + "vitess.io/vitess/go/vt/vtadmin/cache" "vitess.io/vitess/go/vt/vtadmin/cluster" "vitess.io/vitess/go/vt/vtadmin/grpcserver" vtadminhttp "vitess.io/vitess/go/vt/vtadmin/http" @@ -47,6 +48,8 @@ var ( enableRBAC bool disableRBAC bool + cacheRefreshKey string + traceCloser io.Closer = &noopCloser{} rootCmd = &cobra.Command{ @@ -129,6 +132,11 @@ func run(cmd *cobra.Command, args []string) { clusters[i] = cluster } + if cacheRefreshKey == "" { + log.Warningf("no cache-refresh-key set; forcing cache refreshes will not be possible") + } + cache.SetCacheRefreshKey(cacheRefreshKey) + s := vtadmin.NewAPI(clusters, vtadmin.Options{ GRPCOpts: opts, HTTPOpts: httpOpts, @@ -188,6 +196,12 @@ func main() { rootCmd.Flags().BoolVar(&enableRBAC, "rbac", false, "whether to enable RBAC. must be set if not passing --rbac") rootCmd.Flags().BoolVar(&disableRBAC, "no-rbac", false, "whether to disable RBAC. must be set if not passing --no-rbac") + // Global cache flags (N.B. there are also cluster-specific cache flags) + cacheRefreshHelp := "instructs a request to ignore any cached data (if applicable) and refresh the cache;" + + "usable as an HTTP header named 'X-' and as a gRPC metadata key ''\n" + + "Note: any whitespace characters are replaced with hyphens." + rootCmd.Flags().StringVar(&cacheRefreshKey, "cache-refresh-key", "vt-cache-refresh", cacheRefreshHelp) + // glog flags, no better way to do this rootCmd.Flags().AddGoFlag(flag.Lookup("v")) rootCmd.Flags().AddGoFlag(flag.Lookup("logtostderr")) diff --git a/go/vt/vtadmin/api.go b/go/vt/vtadmin/api.go index eec01fdcaca..a2103309829 100644 --- a/go/vt/vtadmin/api.go +++ b/go/vt/vtadmin/api.go @@ -207,6 +207,28 @@ func NewAPI(clusters []*cluster.Cluster, opts Options) *API { return api } +// Close closes all the clusters in an API concurrently. Its primary function is +// to gracefully shutdown cache background goroutines to avoid data races in +// tests, but needs to be exported to be called by those tests. It does not have +// any production use case. +func (api *API) Close() error { + var ( + wg sync.WaitGroup + rec concurrency.AllErrorRecorder + ) + + for _, c := range api.clusters { + wg.Add(1) + go func(c *cluster.Cluster) { + defer wg.Done() + rec.RecordError(c.Close()) + }(c) + } + + wg.Wait() + return rec.Error() +} + // ListenAndServe starts serving this API on the configured Addr (see // grpcserver.Options) until shutdown or irrecoverable error occurs. func (api *API) ListenAndServe() error { diff --git a/go/vt/vtadmin/api_test.go b/go/vt/vtadmin/api_test.go index fb7d7f06f21..03be7ff80ee 100644 --- a/go/vt/vtadmin/api_test.go +++ b/go/vt/vtadmin/api_test.go @@ -29,6 +29,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/grpccommon" @@ -100,6 +101,22 @@ func TestFindSchema(t *testing.T) { }, }, }, + FindAllShardsInKeyspaceResults: map[string]struct { + Response *vtctldatapb.FindAllShardsInKeyspaceResponse + Error error + }{ + "testkeyspace": { + Response: &vtctldatapb.FindAllShardsInKeyspaceResponse{ + Shards: map[string]*vtctldatapb.Shard{ + "-": { + Shard: &topodatapb.Shard{ + IsPrimaryServing: true, + }, + }, + }, + }, + }, + }, }, Tablets: []*vtadminpb.Tablet{ { @@ -109,6 +126,7 @@ func TestFindSchema(t *testing.T) { Uid: 100, }, Keyspace: "testkeyspace", + Shard: "-", }, State: vtadminpb.Tablet_SERVING, }, @@ -531,6 +549,7 @@ func TestFindSchema(t *testing.T) { } api := NewAPI(clusters, Options{}) + defer api.Close() resp, err := api.FindSchema(ctx, tt.req) if tt.shouldErr { @@ -740,6 +759,8 @@ func TestFindSchema(t *testing.T) { ) api := NewAPI([]*cluster.Cluster{c1, c2}, Options{}) + defer api.Close() + schema, err := api.FindSchema(ctx, &vtadminpb.FindSchemaRequest{ Table: "testtable", TableSizeOptions: &vtadminpb.GetSchemaTableSizeOptions{ @@ -774,6 +795,9 @@ func TestFindSchema(t *testing.T) { } if schema != nil { + // Clone so our mutation below doesn't trip the race detector. + schema = proto.Clone(schema).(*vtadminpb.Schema) + for _, td := range schema.TableDefinitions { // Zero these out because they're non-deterministic and also not // relevant to the final result. @@ -1379,7 +1403,6 @@ func TestGetSchema(t *testing.T) { Name: "testtable", }, }, - TableSizes: map[string]*vtadminpb.Schema_TableSize{}, }, shouldErr: false, }, @@ -1529,6 +1552,7 @@ func TestGetSchema(t *testing.T) { Tablets: tt.tablets, }) api := NewAPI([]*cluster.Cluster{c}, Options{}) + defer api.Close() resp, err := api.GetSchema(ctx, tt.req) if tt.shouldErr { @@ -1537,6 +1561,11 @@ func TestGetSchema(t *testing.T) { return } + if resp != nil { + // Clone so our mutation below doesn't trip the race detector. + resp = proto.Clone(resp).(*vtadminpb.Schema) + } + assert.NoError(t, err) assert.Equal(t, tt.expected, resp) }) @@ -1655,6 +1684,8 @@ func TestGetSchema(t *testing.T) { ) api := NewAPI([]*cluster.Cluster{c1, c2}, Options{}) + defer api.Close() + schema, err := api.GetSchema(ctx, &vtadminpb.GetSchemaRequest{ ClusterId: c1.ID, Keyspace: "testkeyspace", @@ -1691,6 +1722,9 @@ func TestGetSchema(t *testing.T) { } if schema != nil { + // Clone so our mutation below doesn't trip the race detector. + schema = proto.Clone(schema).(*vtadminpb.Schema) + for _, td := range schema.TableDefinitions { // Zero these out because they're non-deterministic and also not // relevant to the final result. @@ -2205,6 +2239,7 @@ func TestGetSchemas(t *testing.T) { } api := NewAPI(clusters, Options{}) + defer api.Close() resp, err := api.GetSchemas(ctx, tt.req) require.NoError(t, err) @@ -2425,6 +2460,8 @@ func TestGetSchemas(t *testing.T) { ) api := NewAPI([]*cluster.Cluster{c1, c2}, Options{}) + defer api.Close() + resp, err := api.GetSchemas(ctx, &vtadminpb.GetSchemasRequest{ TableSizeOptions: &vtadminpb.GetSchemaTableSizeOptions{ AggregateSizes: true, @@ -2492,14 +2529,22 @@ func TestGetSchemas(t *testing.T) { } if resp != nil { - for _, schema := range resp.Schemas { + // Clone schemas so our mutations below don't trip the race detector. + schemas := make([]*vtadminpb.Schema, len(resp.Schemas)) + for i, schema := range resp.Schemas { + schema := proto.Clone(schema).(*vtadminpb.Schema) + for _, td := range schema.TableDefinitions { // Zero these out because they're non-deterministic and also not // relevant to the final result. td.RowCount = 0 td.DataLength = 0 } + + schemas[i] = schema } + + resp.Schemas = schemas } assert.NoError(t, err) diff --git a/go/vt/vtadmin/cache/cache.go b/go/vt/vtadmin/cache/cache.go new file mode 100644 index 00000000000..95e92b5d350 --- /dev/null +++ b/go/vt/vtadmin/cache/cache.go @@ -0,0 +1,263 @@ +/* +Copyright 2022 The Vitess Authors. + +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 cache provides a generic key/value cache with support for background +// filling. +package cache + +import ( + "context" + "sync" + "time" + + "github.com/patrickmn/go-cache" + + "vitess.io/vitess/go/vt/log" +) + +// Keyer is the interface cache keys implement to turn themselves into string +// keys. +// +// Note: we define this type rather than using Stringer so users may implement +// that interface for different string representation needs, for example +// providing a human-friendly representation. +type Keyer interface{ Key() string } + +const ( + // DefaultExpiration is used to instruct calls to Add to use the default + // cache expiration. + // See https://pkg.go.dev/github.com/patrickmn/go-cache@v2.1.0+incompatible#pkg-constants. + DefaultExpiration = cache.DefaultExpiration + // NoExpiration is used to create caches that do not expire items by + // default. + // See https://pkg.go.dev/github.com/patrickmn/go-cache@v2.1.0+incompatible#pkg-constants. + NoExpiration = cache.NoExpiration + + // DefaultBackfillEnqueueWaitTime is the default value used for waiting to + // enqueue backfill requests, if a config is passed with a non-positive + // BackfillEnqueueWaitTime. + DefaultBackfillEnqueueWaitTime = time.Millisecond * 50 + // DefaultBackfillRequestTTL is the default value used for how stale of + // backfill requests to still process, if a config is passed with a + // non-positive BackfillRequestTTL. + DefaultBackfillRequestTTL = time.Millisecond * 100 +) + +// Config is the configuration for a cache. +type Config struct { + // DefaultExpiration is how long to keep Values in the cache by default (the + // duration passed to Add takes precedence). Use the sentinel NoExpiration + // to make Values never expire by default. + DefaultExpiration time.Duration `json:"default_expiration"` + // CleanupInterval is how often to remove expired Values from the cache. + CleanupInterval time.Duration `json:"cleanup_interval"` + + // BackfillRequestTTL is how long a backfill request is considered valid. + // If the backfill goroutine encounters a request older than this, it is + // discarded. + BackfillRequestTTL time.Duration `json:"backfill_request_ttl"` + // BackfillRequestDuplicateInterval is how much time must pass before the + // backfill goroutine will re-backfill the same key. It is used to prevent + // multiple callers queuing up too many requests for the same key, when one + // backfill would satisfy all of them. + BackfillRequestDuplicateInterval time.Duration `json:"backfill_request_duplicate_interval"` + // BackfillQueueSize is how many outstanding backfill requests to permit. + // If the queue is full, calls to EnqueueBackfill will return false and + // those requests will be discarded. + BackfillQueueSize int `json:"backfill_queue_size"` + // BackfillEnqueueWaitTime is how long to wait when attempting to enqueue a + // backfill request before giving up. + BackfillEnqueueWaitTime time.Duration `json:"backfill_enqueue_wait_time"` +} + +// Cache is a generic cache supporting background fills. To add things to the +// cache, call Add. To enqueue a background fill, call EnqueueBackfill with a +// Keyer implementation, which will be passed to the fill func provided to New. +// +// For example, to create a schema cache that can backfill full payloads (including +// size aggregation): +// +// var c *cache.Cache[BackfillSchemaRequest, *vtadminpb.Schema] +// c := cache.New(func(ctx context.Context, req BackfillSchemaRequest) (*vtadminpb.Schema, error) { +// // Fetch schema based on fields in `req`. +// // If err is nil, the backfilled schema will be added to the cache. +// return cluster.fetchSchema(ctx, req) +// }) +// +type Cache[Key Keyer, Value any] struct { + cache *cache.Cache + + m sync.Mutex + lastFill map[string]time.Time + + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + backfills chan *backfillRequest[Key] + + fillFunc func(ctx context.Context, k Key) (Value, error) + + cfg Config +} + +// New creates a new cache with the given backfill func. When a request is +// enqueued (via EnqueueBackfill), fillFunc will be called with that request. +func New[Key Keyer, Value any](fillFunc func(ctx context.Context, req Key) (Value, error), cfg Config) *Cache[Key, Value] { + if cfg.BackfillEnqueueWaitTime <= 0 { + log.Warningf("BackfillEnqueueWaitTime (%v) must be positive, defaulting to %v", cfg.BackfillEnqueueWaitTime, DefaultBackfillEnqueueWaitTime) + cfg.BackfillEnqueueWaitTime = DefaultBackfillEnqueueWaitTime + } + + if cfg.BackfillRequestTTL <= 0 { + log.Warningf("BackfillRequestTTL (%v) must be positive, defaulting to %v", cfg.BackfillRequestTTL, DefaultBackfillRequestTTL) + cfg.BackfillRequestTTL = DefaultBackfillRequestTTL + } + + c := &Cache[Key, Value]{ + cache: cache.New(cfg.DefaultExpiration, cfg.CleanupInterval), + lastFill: map[string]time.Time{}, + backfills: make(chan *backfillRequest[Key], cfg.BackfillQueueSize), + fillFunc: fillFunc, + cfg: cfg, + } + + c.ctx, c.cancel = context.WithCancel(context.Background()) + + c.wg.Add(1) + go c.backfill() // TODO: consider allowing N backfill threads to run, configurable + + return c +} + +// Add adds a (key, value) to the cache directly, following the semantics of +// (github.com/patrickmn/go-cache).Cache.Add. +func (c *Cache[Key, Value]) Add(key Key, val Value, d time.Duration) error { + return c.add(key.Key(), val, d) +} + +func (c *Cache[Key, Value]) add(key string, val Value, d time.Duration) error { + c.m.Lock() + // Record the time we last cached this key, to check against + c.lastFill[key] = time.Now().UTC() + c.m.Unlock() + + // Then cache the actual value. + return c.cache.Add(key, val, d) +} + +// Get returns the Value stored for the key, if present in the cache. If the key +// is not cached, the zero value for the given type is returned, along with a +// boolean to indicated presence/absence. +func (c *Cache[Key, Value]) Get(key Key) (Value, bool) { + v, exp, ok := c.cache.GetWithExpiration(key.Key()) + if !ok || (!exp.IsZero() && exp.Before(time.Now())) { + var zero Value + return zero, false + } + + return v.(Value), ok +} + +// EnqueueBackfill submits a request to the backfill queue. +func (c *Cache[Key, Value]) EnqueueBackfill(k Key) bool { + req := &backfillRequest[Key]{ + k: k, + requestedAt: time.Now().UTC(), + } + + select { + case c.backfills <- req: + return true + case <-time.After(c.cfg.BackfillEnqueueWaitTime): + return false + } +} + +// Close closes the backfill goroutine, effectively rendering this cache +// unusable for further background fills. +func (c *Cache[Key, Value]) Close() error { + c.cancel() + c.wg.Wait() + return nil +} + +type backfillRequest[Key Keyer] struct { + k Key + requestedAt time.Time +} + +func (c *Cache[Key, Value]) backfill() { + defer c.wg.Done() + + for { + var req *backfillRequest[Key] + select { + case <-c.ctx.Done(): + return + case req = <-c.backfills: + } + + if req.requestedAt.Add(c.cfg.BackfillRequestTTL).Before(time.Now()) { + // We took too long to get to this request, per config options. + log.Warningf("backfill for %s requested at %s; discarding due to exceeding TTL (%s)", req.k.Key(), req.requestedAt, c.cfg.BackfillRequestTTL) + continue + } + + key := req.k.Key() + + c.m.Lock() + if t, ok := c.lastFill[key]; ok { + if !t.IsZero() && t.Add(c.cfg.BackfillRequestDuplicateInterval).After(time.Now()) { + // We recently added a value for this key to the cache, either via + // another backfill request, or directly via a call to Add. + log.Infof("filled cache for %s less than %s ago (at %s)", key, c.cfg.BackfillRequestDuplicateInterval, t.UTC()) + c.m.Unlock() + continue + } + + // NOTE: In the strictest sense, we would `delete(lastFill, key)` + // here. However, we're about to fill that key again, so we'd be + // immediately re-adding the key we just deleted from `lastFill`. + // + // We do *not* apply the same treatment to the actual Value cache, + // because go-cache is running a background thread to periodically + // clean up expired entries. + } + c.m.Unlock() + + val, err := c.fillFunc(c.ctx, req.k) + if err != nil { + log.Errorf("backfill failed for key %s: %s", key, err) + // TODO: consider re-requesting with a retry-counter paired with a config to give up after N attempts + continue + } + + // Finally, store the value. + if err := c.add(key, val, cache.DefaultExpiration); err != nil { + log.Warningf("failed to add (%s, %+v) to cache: %s", key, val, err) + } + } +} + +// Debug implements debug.Debuggable for Cache. +func (c *Cache[Key, Value]) Debug() map[string]any { + return map[string]any{ + "size": c.cache.ItemCount(), // NOTE: this may include expired items that have not been evicted yet. + "config": c.cfg, + "backfill_queue_length": len(c.backfills), + "closed": c.ctx.Err() != nil, + } +} diff --git a/go/vt/vtadmin/cache/cache_test.go b/go/vt/vtadmin/cache/cache_test.go new file mode 100644 index 00000000000..93a6898db5d --- /dev/null +++ b/go/vt/vtadmin/cache/cache_test.go @@ -0,0 +1,193 @@ +/* +Copyright 2022 The Vitess Authors. + +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 cache_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "vitess.io/vitess/go/vt/vtadmin/cache" +) + +type testkey string + +func (k testkey) Key() string { return string(k) } + +func TestBackfillDuplicates(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg cache.Config + enqueueInterval time.Duration + enqueueCount int + expectedCallCount int + assertionMsg string + }{ + { + name: "duplicates too close together", + cfg: cache.Config{ + BackfillRequestTTL: time.Hour, + BackfillRequestDuplicateInterval: time.Second, + }, + enqueueInterval: 0, // no sleep between enqueues + enqueueCount: 2, + expectedCallCount: 1, + assertionMsg: "fillFunc should only have been called once for rapid duplicate requests", + }, + { + name: "duplicates with enough time passed", + cfg: cache.Config{ + BackfillRequestTTL: time.Hour, + BackfillRequestDuplicateInterval: time.Millisecond, + }, + enqueueInterval: time.Millisecond * 5, // sleep longer than the duplicate interval + enqueueCount: 2, + expectedCallCount: 2, + assertionMsg: "fillFunc should have been called for duplicate requests with enough time between them", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var callcount int + c := cache.New(func(ctx context.Context, req testkey) (any, error) { + callcount++ + return nil, nil + }, tt.cfg) + + key := testkey("testkey") + for i := 0; i < tt.enqueueCount; i++ { + if !c.EnqueueBackfill(key) { + assert.Fail(t, "failed to enqueue backfill for key %s", key) + } + + time.Sleep(tt.enqueueInterval) + } + + c.Close() + + assert.Equal(t, tt.expectedCallCount, callcount, tt.assertionMsg) + }) + } +} + +func TestBackfillTTL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg cache.Config + fillSleep time.Duration + enqueueCount int + expectedCallCount int + assertionMsg string + }{ + { + name: "within backfill ttl", + cfg: cache.Config{ + BackfillRequestTTL: time.Millisecond * 50, + BackfillRequestDuplicateInterval: time.Second, + BackfillEnqueueWaitTime: time.Second, + }, + fillSleep: time.Millisecond * 10, + enqueueCount: 2, + expectedCallCount: 1, + assertionMsg: "both requests are within TTL and should have been filled", + }, + { + name: "backfill ttl exceeded", + cfg: cache.Config{ + BackfillRequestTTL: time.Millisecond * 10, + BackfillRequestDuplicateInterval: time.Millisecond, + BackfillEnqueueWaitTime: time.Second, + }, + fillSleep: time.Millisecond * 50, + enqueueCount: 2, + expectedCallCount: 1, + assertionMsg: "second request exceeds TTL and should not have been filled", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var callcount int + c := cache.New(func(ctx context.Context, req testkey) (any, error) { + time.Sleep(tt.fillSleep) // make fills take time so that requests in the queue can exceed ttl + callcount++ + return nil, nil + }, tt.cfg) + + key := testkey("testkey") + for i := 0; i < tt.enqueueCount; i++ { + if !c.EnqueueBackfill(key) { + assert.Fail(t, "failed to enqueue backfill for key %s", key) + } + } + + time.Sleep(tt.fillSleep * time.Duration(tt.enqueueCount)) + c.Close() + + assert.Equal(t, tt.expectedCallCount, callcount, tt.assertionMsg) + }) + } +} + +type intkey int + +func (k intkey) Key() string { return fmt.Sprintf("%d", int(k)) } + +func TestEnqueueBackfillTimeout(t *testing.T) { + t.Parallel() + + var callcount int + c := cache.New(func(ctx context.Context, req intkey) (any, error) { + time.Sleep(time.Millisecond * 50) // make fills take time so that the second enqueue exceeds WaitTimeout + callcount++ + return nil, nil + }, cache.Config{ + BackfillEnqueueWaitTime: time.Millisecond * 10, + }) + + var enqueues = []struct { + shouldFail bool + msg string + }{ + { + shouldFail: false, + msg: "not exceed", + }, + { + shouldFail: true, + msg: "exceed", + }, + } + for i, q := range enqueues { + ok := c.EnqueueBackfill(intkey(i)) + assert.Equal(t, q.shouldFail, !ok, "enqueue should %s wait timeout", q.msg) + } +} diff --git a/go/vt/vtadmin/cache/refresh.go b/go/vt/vtadmin/cache/refresh.go new file mode 100644 index 00000000000..1e37d5ca41d --- /dev/null +++ b/go/vt/vtadmin/cache/refresh.go @@ -0,0 +1,100 @@ +/* +Copyright 2022 The Vitess Authors. + +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 cache + +import ( + "context" + "fmt" + "net/http" + "strconv" + "strings" + + "google.golang.org/grpc/metadata" + + "vitess.io/vitess/go/vt/log" +) + +var ( + cacheRefreshHeader string + cacheRefreshGRPCMetadataKey string +) + +// SetCacheRefreshKey sets the global HTTP header and gRPC metadata keys for +// requests to force a cache refresh. Any whitespace characters are replaced +// with hyphens. +// +// It is not threadsafe, and should be called only at startup or in an init +// function. +func SetCacheRefreshKey(k string) { + l := strings.ToLower(k) + l = strings.ReplaceAll(l, " ", "-") + cacheRefreshHeader, cacheRefreshGRPCMetadataKey = fmt.Sprintf("x-%s", l), l +} + +// NewIncomingRefreshContext returns an incoming gRPC context with metadata +// set to signal a cache refresh. +func NewIncomingRefreshContext(ctx context.Context) context.Context { + md := metadata.Pairs(cacheRefreshGRPCMetadataKey, "true") + return metadata.NewIncomingContext(ctx, md) +} + +// ShouldRefreshFromIncomingContext returns true if the gRPC metadata in the +// incoming context signals a cache refresh was requested. +func ShouldRefreshFromIncomingContext(ctx context.Context) bool { + if cacheRefreshGRPCMetadataKey == "" { + return false + } + + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return false + } + + vals := md.Get(cacheRefreshGRPCMetadataKey) + if len(vals) == 0 { + return false + } + + shouldRefresh, err := strconv.ParseBool(vals[0]) + if err != nil { + log.Warningf("failed to parse %s metadata key as bool: %s", cacheRefreshGRPCMetadataKey, err) + return false + } + + return shouldRefresh +} + +// ShouldRefreshFromRequest returns true if the HTTP request headers signal a +// cache refresh was requested. +func ShouldRefreshFromRequest(r *http.Request) bool { + if cacheRefreshHeader == "" { + return false + } + + h := r.Header.Get(cacheRefreshHeader) + if h == "" { + return false + } + + shouldRefresh, err := strconv.ParseBool(h) + if err != nil { + log.Warningf("failed to parse %s header as bool: %s", cacheRefreshHeader, err) + return false + } + + return shouldRefresh +} diff --git a/go/vt/vtadmin/cache/refresh_test.go b/go/vt/vtadmin/cache/refresh_test.go new file mode 100644 index 00000000000..c12bb63ad6a --- /dev/null +++ b/go/vt/vtadmin/cache/refresh_test.go @@ -0,0 +1,130 @@ +/* +Copyright 2022 The Vitess Authors. + +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 cache_test + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/metadata" + + "vitess.io/vitess/go/vt/vtadmin/cache" +) + +const refreshKey = "cache_test" + +func init() { + cache.SetCacheRefreshKey(refreshKey) +} + +func TestShouldRefreshFromIncomingContext(t *testing.T) { + t.Parallel() + + t.Run("true", func(t *testing.T) { + t.Parallel() + + ctx := cache.NewIncomingRefreshContext(context.Background()) + assert.True(t, cache.ShouldRefreshFromIncomingContext(ctx), "incoming context with metadata set should signal refresh") + }) + + t.Run("false", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ctx context.Context + }{ + { + name: "no metadata in context", + ctx: context.Background(), + }, + { + name: "no cache key in metadata", + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("key", "val")), + }, + { + name: "cache key set to non-bool", + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(refreshKey, "this is not a boolean")), + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + assert.False(t, cache.ShouldRefreshFromIncomingContext(tt.ctx), "incoming context should not signal refresh") + }) + } + }) +} + +func TestShouldRefreshFromRequest(t *testing.T) { + t.Parallel() + + t.Run("true", func(t *testing.T) { + t.Parallel() + + r, _ := http.NewRequest("", "", nil) + r.Header.Add("x-"+refreshKey, "true") + + assert.True(t, cache.ShouldRefreshFromRequest(r), "request header should signal cache refresh") + }) + + t.Run("false", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + headers map[string]string + }{ + { + name: "no headers set", + headers: nil, + }, + { + name: "no cache refresh header", + headers: map[string]string{ + "x-forwarded-for": "whatever", + "content-type": "application/json", + }, + }, + { + name: "cache refresh set to non-boolean", + headers: map[string]string{ + "x-" + refreshKey: "this is not a boolean", + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + r, _ := http.NewRequest("", "", nil) + for header, val := range tt.headers { + r.Header.Add(header, val) + } + + assert.False(t, cache.ShouldRefreshFromRequest(r), "request headers should not signal cache refresh") + }) + } + }) +} diff --git a/go/vt/vtadmin/cluster/cluster.go b/go/vt/vtadmin/cluster/cluster.go index 8a070f30f26..50b49544b54 100644 --- a/go/vt/vtadmin/cluster/cluster.go +++ b/go/vt/vtadmin/cluster/cluster.go @@ -40,7 +40,9 @@ import ( "vitess.io/vitess/go/vt/log" "vitess.io/vitess/go/vt/logutil" "vitess.io/vitess/go/vt/topo/topoproto" + "vitess.io/vitess/go/vt/vtadmin/cache" "vitess.io/vitess/go/vt/vtadmin/cluster/discovery" + "vitess.io/vitess/go/vt/vtadmin/cluster/internal/caches/schemacache" "vitess.io/vitess/go/vt/vtadmin/debug" "vitess.io/vitess/go/vt/vtadmin/errors" "vitess.io/vitess/go/vt/vtadmin/vtadminproto" @@ -80,6 +82,17 @@ type Cluster struct { emergencyReparentPool *pools.RPCPool // ERS-only reparentPool *pools.RPCPool // PRS-only + // schemaCache caches schema(s) for different GetSchema(s) requests. + // + // - if we call GetSchema, then getSchemaCacheRequest.Keyspace will be + // non-empty and the cached schemas slice will contain exactly one element, + // namely for that keyspace's schema. + // - if we call GetSchemas, then getSchemaCacheRequest == "", and the cached + // schemas slice will contain one element per keyspace* in the cluster + // *: at the time it was cached; if keyspaces were created/destroyed in + // the interim, we won't pick that up until something refreshes the cache. + schemaCache *cache.Cache[schemacache.Key, []*vtadminpb.Schema] + cfg Config } @@ -150,9 +163,51 @@ func New(ctx context.Context, cfg Config) (*Cluster, error) { cluster.emergencyReparentPool = cfg.EmergencyReparentPoolConfig.NewRWPool() cluster.reparentPool = cfg.ReparentPoolConfig.NewRWPool() + if cluster.cfg.SchemaCacheConfig == nil { + cluster.cfg.SchemaCacheConfig = &cache.Config{} + } + cluster.schemaCache = cache.New(func(ctx context.Context, key schemacache.Key) ([]*vtadminpb.Schema, error) { + // TODO: make a private method to separate the fetching bits from the cache bits + if key.Keyspace == "" { + return cluster.GetSchemas(ctx, GetSchemaOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + }, + TableSizeOptions: &vtadminpb.GetSchemaTableSizeOptions{ + AggregateSizes: true, + IncludeNonServingShards: key.IncludeNonServingShards, + }, + isBackfill: true, + }) + } + + schema, err := cluster.GetSchema(ctx, key.Keyspace, GetSchemaOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + }, + TableSizeOptions: &vtadminpb.GetSchemaTableSizeOptions{ + AggregateSizes: true, + IncludeNonServingShards: key.IncludeNonServingShards, + }, + isBackfill: true, + }) + if err != nil { + return nil, err + } + + return []*vtadminpb.Schema{schema}, nil + }, *cluster.cfg.SchemaCacheConfig) + return cluster, nil } +// Close closes a cluster. Its primary function is to gracefully shutdown any +// background cache goroutines to avoid data races in tests. It does not have +// any production use case, but needs to be exported. +func (c *Cluster) Close() error { + return c.schemaCache.Close() +} + // ToProto returns a value-copy protobuf equivalent of the cluster. func (c Cluster) ToProto() *vtadminpb.Cluster { return &vtadminpb.Cluster{ @@ -1161,6 +1216,8 @@ type GetSchemaOptions struct { // (described above) to find one SERVING tablet for each shard in the // keyspace, skipping any non-serving shards in the keyspace. TableSizeOptions *vtadminpb.GetSchemaTableSizeOptions + + isBackfill bool } // GetSchema returns the schema for a given keyspace. GetSchema has a few @@ -1208,6 +1265,24 @@ func (c *Cluster) GetSchema(ctx context.Context, keyspace string, opts GetSchema span.Annotate("keyspace", keyspace) annotateGetSchemaRequest(opts.BaseRequest, span) vtadminproto.AnnotateSpanWithGetSchemaTableSizeOptions(opts.TableSizeOptions, span) + span.Annotate("is_backfill", opts.isBackfill) + + key := schemacache.Key{ + ClusterID: c.ID, + Keyspace: keyspace, + IncludeNonServingShards: opts.TableSizeOptions.IncludeNonServingShards, + } + if !(opts.isBackfill || cache.ShouldRefreshFromIncomingContext(ctx)) { + schema, ok, err := schemacache.LoadOne(c.schemaCache, key, schemacache.LoadOptions{ + BaseRequest: opts.BaseRequest, + AggregateSizes: opts.TableSizeOptions.AggregateSizes, + }) + + span.Annotate("cache_hit", ok) + if ok { + return schema, err + } + } // Fetch all tablets for the keyspace. tablets, err := c.FindTablets(ctx, func(tablet *vtadminpb.Tablet) bool { @@ -1222,7 +1297,17 @@ func (c *Cluster) GetSchema(ctx context.Context, keyspace string, opts GetSchema return nil, err } - return c.getSchemaFromTablets(ctx, keyspace, tabletsToQuery, opts) + schema, err := c.getSchemaFromTablets(ctx, keyspace, tabletsToQuery, opts) + if err != nil { + return nil, err + } + + go schemacache.AddOrBackfill(c.schemaCache, []*vtadminpb.Schema{schema}, key, cache.DefaultExpiration, schemacache.LoadOptions{ + BaseRequest: opts.BaseRequest, + AggregateSizes: opts.TableSizeOptions.AggregateSizes, + }) + + return schema, nil } // GetSchemas returns all of the schemas across all keyspaces in the cluster. @@ -1249,6 +1334,24 @@ func (c *Cluster) GetSchemas(ctx context.Context, opts GetSchemaOptions) ([]*vta AnnotateSpan(c, span) annotateGetSchemaRequest(opts.BaseRequest, span) vtadminproto.AnnotateSpanWithGetSchemaTableSizeOptions(opts.TableSizeOptions, span) + span.Annotate("is_backfill", opts.isBackfill) + + key := schemacache.Key{ + ClusterID: c.ID, + Keyspace: "", + IncludeNonServingShards: opts.TableSizeOptions.IncludeNonServingShards, + } + if !(opts.isBackfill || cache.ShouldRefreshFromIncomingContext(ctx)) { + schemas, ok, err := schemacache.LoadAll(c.schemaCache, key, schemacache.LoadOptions{ + BaseRequest: opts.BaseRequest, + AggregateSizes: opts.TableSizeOptions.AggregateSizes, + }) + + span.Annotate("cache_hit", ok) + if ok { + return schemas, err + } + } var ( m sync.Mutex @@ -1357,6 +1460,11 @@ func (c *Cluster) GetSchemas(ctx context.Context, opts GetSchemaOptions) ([]*vta return nil, rec.Error() } + go schemacache.AddOrBackfill(c.schemaCache, schemas, key, cache.DefaultExpiration, schemacache.LoadOptions{ + BaseRequest: opts.BaseRequest, + AggregateSizes: opts.TableSizeOptions.AggregateSizes, + }) + return schemas, nil } @@ -2198,6 +2306,9 @@ func (c *Cluster) Debug() map[string]any { "emergency_reparent_pool": json.RawMessage(c.emergencyReparentPool.StatsJSON()), "reparent_pool": json.RawMessage(c.reparentPool.StatsJSON()), }, + "caches": map[string]any{ + "schemas": c.schemaCache.Debug(), + }, } if vtsql, ok := c.DB.(debug.Debuggable); ok { diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index e388ae16585..3c2d8ae990e 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -27,6 +27,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" "k8s.io/apimachinery/pkg/util/sets" "vitess.io/vitess/go/mysql" @@ -2695,6 +2696,9 @@ func TestGetSchema(t *testing.T) { return } + // Clone so our mutation below doesn't trip the race detector. + schema = proto.Clone(schema).(*vtadminpb.Schema) + if schema.TableDefinitions != nil { // For simplicity, we're going to assert only on the state // of the aggregated sizes (in schema.TableSizes), since the diff --git a/go/vt/vtadmin/cluster/config.go b/go/vt/vtadmin/cluster/config.go index 9b2e7fdfdec..c27e6975d2f 100644 --- a/go/vt/vtadmin/cluster/config.go +++ b/go/vt/vtadmin/cluster/config.go @@ -29,6 +29,7 @@ import ( "vitess.io/vitess/go/pools" "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/vtadmin/cache" "vitess.io/vitess/go/vt/vtadmin/errors" "vitess.io/vitess/go/vt/vtadmin/vtctldclient" "vitess.io/vitess/go/vt/vtadmin/vtsql" @@ -74,6 +75,8 @@ type Config struct { // RW RPCPool. ReparentPoolConfig *RPCPoolConfig + SchemaCacheConfig *cache.Config + vtctldConfigOpts []vtctldclient.ConfigOption vtsqlConfigOpts []vtsql.ConfigOption } @@ -185,6 +188,11 @@ func (cfg *Config) MarshalJSON() ([]byte, error) { Size: DefaultReadPoolSize, WaitTimeout: DefaultReadPoolWaitTimeout, } + defaultCacheConfig := &cache.Config{ + BackfillRequestTTL: cache.DefaultBackfillRequestTTL, + BackfillQueueSize: 10, + BackfillEnqueueWaitTime: cache.DefaultBackfillEnqueueWaitTime, + } tmp := struct { ID string `json:"id"` @@ -203,6 +211,8 @@ func (cfg *Config) MarshalJSON() ([]byte, error) { EmergencyReparentPoolConfig *RPCPoolConfig `json:"emergency_reparent_pool_config"` ReparentPoolConfig *RPCPoolConfig `json:"reparent_pool_config"` + + SchemaCacheConfig *cache.Config `json:"schema_cache_config"` }{ ID: cfg.ID, Name: cfg.Name, @@ -217,6 +227,7 @@ func (cfg *Config) MarshalJSON() ([]byte, error) { WorkflowReadPoolConfig: defaultReadPoolConfig.merge(cfg.WorkflowReadPoolConfig), EmergencyReparentPoolConfig: defaultRWPoolConfig.merge(cfg.EmergencyReparentPoolConfig), ReparentPoolConfig: defaultRWPoolConfig.merge(cfg.ReparentPoolConfig), + SchemaCacheConfig: mergeCacheConfigs(defaultCacheConfig, cfg.SchemaCacheConfig), } return json.Marshal(&tmp) @@ -240,6 +251,7 @@ func (cfg Config) Merge(override Config) Config { WorkflowReadPoolConfig: cfg.WorkflowReadPoolConfig.merge(override.WorkflowReadPoolConfig), EmergencyReparentPoolConfig: cfg.EmergencyReparentPoolConfig.merge(override.EmergencyReparentPoolConfig), ReparentPoolConfig: cfg.ReparentPoolConfig.merge(override.ReparentPoolConfig), + SchemaCacheConfig: mergeCacheConfigs(cfg.SchemaCacheConfig, override.SchemaCacheConfig), } if override.ID != "" { @@ -278,6 +290,81 @@ func mergeStringMap(base map[string]string, override map[string]string) { } } +func mergeCacheConfigs(base, override *cache.Config) *cache.Config { + if base == nil && override == nil { + return nil + } + + merged := &cache.Config{ + DefaultExpiration: -1, + CleanupInterval: -1, + BackfillRequestTTL: -1, + BackfillRequestDuplicateInterval: -1, + BackfillQueueSize: -1, + BackfillEnqueueWaitTime: -1, + } + + for _, c := range []*cache.Config{base, override} { + if c != nil { + if c.DefaultExpiration >= 0 { + merged.DefaultExpiration = c.DefaultExpiration + } + + if c.CleanupInterval >= 0 { + merged.CleanupInterval = c.CleanupInterval + } + + if c.BackfillRequestTTL >= 0 { + merged.BackfillRequestTTL = c.BackfillRequestTTL + } + + if c.BackfillRequestDuplicateInterval >= 0 { + merged.BackfillRequestDuplicateInterval = c.BackfillRequestDuplicateInterval + } + + if c.BackfillQueueSize >= 0 { + merged.BackfillQueueSize = c.BackfillQueueSize + } + + if c.BackfillEnqueueWaitTime >= 0 { + merged.BackfillEnqueueWaitTime = c.BackfillEnqueueWaitTime + } + } + } + + return merged +} + +func parseCacheConfigFlag(cfg *cache.Config, name, val string) (err error) { + switch name { + case "default-expiration": + cfg.DefaultExpiration, err = time.ParseDuration(val) + case "cleanup-interval": + cfg.CleanupInterval, err = time.ParseDuration(val) + case "backfill-request-ttl": + cfg.BackfillRequestTTL, err = time.ParseDuration(val) + case "backfill-request-duplicate-interval": + cfg.BackfillRequestDuplicateInterval, err = time.ParseDuration(val) + case "backfill-queue-size": + size, err := strconv.ParseInt(val, 10, 64) + if err != nil { + return err + } + + if size < 0 { + return fmt.Errorf("%w: backfill queue size must be non-negative; got %d", strconv.ErrRange, size) + } + + cfg.BackfillQueueSize = int(size) + case "backfill-enqueue-wait-time": + cfg.BackfillEnqueueWaitTime, err = time.ParseDuration(val) + default: + return errors.ErrNoFlag + } + + return err +} + // RPCPoolConfig holds configuration options for creating RPCPools. type RPCPoolConfig struct { Size int `json:"size"` diff --git a/go/vt/vtadmin/cluster/flags.go b/go/vt/vtadmin/cluster/flags.go index b5c203c27b3..891c6d10c44 100644 --- a/go/vt/vtadmin/cluster/flags.go +++ b/go/vt/vtadmin/cluster/flags.go @@ -22,6 +22,7 @@ import ( "strings" "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/vtadmin/cache" ) // FlagsByImpl groups a set of flags by discovery implementation. Its mapping is @@ -237,6 +238,21 @@ func parseOne(cfg *Config, name string, val string) error { if err := cfg.ReparentPoolConfig.parseFlag(strings.TrimPrefix(name, "reparent-pool-"), val); err != nil { return fmt.Errorf("error parsing %s: %w", name, err) } + case strings.HasPrefix(name, "schema-cache-"): + if cfg.SchemaCacheConfig == nil { + cfg.SchemaCacheConfig = &cache.Config{ + DefaultExpiration: -1, + CleanupInterval: -1, + BackfillRequestTTL: -1, + BackfillRequestDuplicateInterval: -1, + BackfillQueueSize: -1, + BackfillEnqueueWaitTime: -1, + } + } + + if err := parseCacheConfigFlag(cfg.SchemaCacheConfig, strings.TrimPrefix(name, "schema-cache-"), val); err != nil { + return fmt.Errorf("error parsing %s: %w", name, err) + } default: match := discoveryFlagRegexp.FindStringSubmatch(name) if match == nil { diff --git a/go/vt/vtadmin/cluster/internal/caches/schemacache/cache.go b/go/vt/vtadmin/cluster/internal/caches/schemacache/cache.go new file mode 100644 index 00000000000..8a7612e89a7 --- /dev/null +++ b/go/vt/vtadmin/cluster/internal/caches/schemacache/cache.go @@ -0,0 +1,218 @@ +/* +Copyright 2022 The Vitess Authors. + +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 schemacache provides wrapper functions for interacting with +// instances of the generic (vtadmin/cache).Cache that store schemas. +package schemacache + +import ( + "fmt" + "time" + + "google.golang.org/protobuf/proto" + + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/mysqlctl/tmutils" + "vitess.io/vitess/go/vt/vtadmin/cache" + "vitess.io/vitess/go/vt/vterrors" + + tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata" + vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin" + vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata" + "vitess.io/vitess/go/vt/proto/vtrpc" +) + +// Key is the cache key for vtadmin's schema caches. +type Key struct { + ClusterID string + Keyspace string + IncludeNonServingShards bool +} + +// Key implements cache.Keyer for Key. +func (k Key) Key() string { + return fmt.Sprintf("%s/%s/%v", k.ClusterID, k.Keyspace, k.IncludeNonServingShards) +} + +type schemaCache = cache.Cache[Key, []*vtadminpb.Schema] // define a type alias for brevity within this package. + +// AddOrBackfill takes a (schemas, key, duration) triple (see cache.Add) and a +// LoadOptions and either adds the schemas to the cache directly for the key, or +// enqueues a backfill for the key. Which of these occurs depends on the given +// LoadOptions. If the options, when passed to Load(All|One) would result in a +// full payload of the cached data, then the schemas are added to the cache +// directly. Otherwise, data would be missing for future Load(All|One) calls, +// so we enqueue a backfill to fetch the full payload and add that to the cache. +// +// Callers are expected to call this function in the background, so it returns +// nothing and only logs warnings on failures. +func AddOrBackfill(c *schemaCache, schemas []*vtadminpb.Schema, key Key, d time.Duration, opts LoadOptions) { + if opts.isFullPayload() { + if err := c.Add(key, schemas, d); err != nil { + log.Warningf("failed to add schema to cache for %+v: %s", key, err) + } + } else { + if !c.EnqueueBackfill(key) { + log.Warningf("failed to enqueue backfill for schema cache %+v", key) + } + } +} + +// LoadOptions is the set of options used by Load(All|One) to filter down fully- +// cached schema payloads. +type LoadOptions struct { + BaseRequest *vtctldatapb.GetSchemaRequest + AggregateSizes bool + filter *tmutils.TableFilter +} + +// isFullPayload returns true if the schema(s) returned by calls to +// Load(All|One) from these options represents the full cached payload. if +// false, Load(All|One) will return only a subset of what is in the cache, +// meaning that AddOrBackfill for these options should enqueue a backfill to +// fetch the full payload. +func (opts LoadOptions) isFullPayload() bool { + if opts.BaseRequest == nil { + return false + } + + if len(opts.BaseRequest.Tables) > 0 { + return false + } + + if len(opts.BaseRequest.ExcludeTables) > 0 { + return false + } + + if !opts.BaseRequest.IncludeViews { + return false + } + + if opts.BaseRequest.TableNamesOnly { + return false + } + + if opts.BaseRequest.TableSizesOnly { + return false + } + + if !opts.AggregateSizes { + return false + } + + return true +} + +// LoadAll returns the slice of all schemas cached for the given key, filtering +// down based on LoadOptions. +// +// The boolean return value will be false if there is nothing in the cache for +// key. +func LoadAll(c *schemaCache, key Key, opts LoadOptions) (schemas []*vtadminpb.Schema, ok bool, err error) { + cachedSchemas, ok := c.Get(key) + if !ok { + return nil, false, nil + } + + opts.filter, err = tmutils.NewTableFilter(opts.BaseRequest.Tables, opts.BaseRequest.ExcludeTables, opts.BaseRequest.IncludeViews) + if err != nil { + // Note that there's no point in falling back to a full fetch in this + // case; this exact code gets executed on the tabletserver side, too, + // so we would just fail on the same issue _after_ paying the cost to + // wait for a roundtrip to the remote tablet. + return nil, ok, err + } + + schemas = make([]*vtadminpb.Schema, 0, len(cachedSchemas)) + for _, schema := range cachedSchemas { + if key.Keyspace != "" && schema.Keyspace != key.Keyspace { + continue + } + + schemas = append(schemas, loadSchema(schema, opts)) + } + + return schemas, ok, nil +} + +// LoadOne loads a single schema for the given key, filtering down based on the +// LoadOptions. If there is not exactly one schema for the key, an error is +// returned. +// +// The boolean return value will be false if there is nothing in the cache for +// key. +func LoadOne(c *schemaCache, key Key, opts LoadOptions) (schema *vtadminpb.Schema, found bool, err error) { + schemas, found, err := LoadAll(c, key, opts) + if err != nil { + return nil, found, err + } + + if !found { + return nil, found, nil + } + + if len(schemas) != 1 { + err = vterrors.Errorf(vtrpc.Code_INTERNAL, "impossible: cache should have exactly 1 schema for request %+v (have %d)", key, len(schemas)) + log.Errorf(err.Error()) + return nil, found, err + } + + return schemas[0], found, nil +} + +func loadSchema(cachedSchema *vtadminpb.Schema, opts LoadOptions) *vtadminpb.Schema { + schema := proto.Clone(cachedSchema).(*vtadminpb.Schema) + + if !opts.AggregateSizes { + schema.TableSizes = nil + } + + tables := make([]*tabletmanagerdatapb.TableDefinition, 0, len(schema.TableDefinitions)) + tableSizes := make(map[string]*vtadminpb.Schema_TableSize, len(schema.TableSizes)) + + for _, td := range schema.TableDefinitions { + if !opts.filter.Includes(td.Name, td.Type) { + continue + } + + if opts.AggregateSizes { + tableSizes[td.Name] = schema.TableSizes[td.Name] + } + + if opts.BaseRequest.TableNamesOnly { + tables = append(tables, &tabletmanagerdatapb.TableDefinition{ + Name: td.Name, + }) + continue + } + + if opts.BaseRequest.TableSizesOnly { + tables = append(tables, &tabletmanagerdatapb.TableDefinition{ + Name: td.Name, + DataLength: td.DataLength, + RowCount: td.RowCount, + }) + continue + } + + tables = append(tables, td) + } + + schema.TableDefinitions = tables + schema.TableSizes = tableSizes + + return schema +} diff --git a/go/vt/vtadmin/cluster/internal/caches/schemacache/cache_test.go b/go/vt/vtadmin/cluster/internal/caches/schemacache/cache_test.go new file mode 100644 index 00000000000..aae706bbad6 --- /dev/null +++ b/go/vt/vtadmin/cluster/internal/caches/schemacache/cache_test.go @@ -0,0 +1,131 @@ +/* +Copyright 2022 The Vitess Authors. + +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 schemacache + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata" +) + +func TestLoadOptions(t *testing.T) { + t.Parallel() + + t.Run("isFullPayload", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts LoadOptions + isFullPayload bool + }{ + { + name: "full payload", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + }, + AggregateSizes: true, + }, + isFullPayload: true, + }, + { + name: "tables", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + Tables: []string{"t1"}, + IncludeViews: true, + }, + AggregateSizes: true, + }, + isFullPayload: false, + }, + { + name: "exclude tables", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + ExcludeTables: []string{"t1"}, + IncludeViews: true, + }, + AggregateSizes: true, + }, + isFullPayload: false, + }, + { + name: "no views", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: false, + }, + AggregateSizes: true, + }, + isFullPayload: false, + }, + { + name: "names only", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + TableNamesOnly: true, + }, + AggregateSizes: true, + }, + isFullPayload: false, + }, + { + name: "sizes only", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + TableSizesOnly: true, + }, + AggregateSizes: true, + }, + isFullPayload: false, + }, + { + name: "no size aggregation", + opts: LoadOptions{ + BaseRequest: &vtctldatapb.GetSchemaRequest{ + IncludeViews: true, + }, + AggregateSizes: false, + }, + isFullPayload: false, + }, + { + name: "missing request", + opts: LoadOptions{ + BaseRequest: nil, + AggregateSizes: true, + }, + isFullPayload: false, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + assert.Equal(t, tt.isFullPayload, tt.opts.isFullPayload(), "%+v", tt.opts) + }) + } + }) +} diff --git a/go/vt/vtadmin/http/api.go b/go/vt/vtadmin/http/api.go index 9c3f53481a8..4d5d54b3afb 100644 --- a/go/vt/vtadmin/http/api.go +++ b/go/vt/vtadmin/http/api.go @@ -21,6 +21,7 @@ import ( "net/http" "vitess.io/vitess/go/trace" + "vitess.io/vitess/go/vt/vtadmin/cache" "vitess.io/vitess/go/vt/vtadmin/rbac" vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin" @@ -79,6 +80,10 @@ func (api *API) Adapt(handler VTAdminHandler) http.HandlerFunc { ctx = rbac.NewContext(ctx, actor) } + if cache.ShouldRefreshFromRequest(r) { + ctx = cache.NewIncomingRefreshContext(ctx) + } + handler(ctx, Request{r}, api).Write(w) } }