[vtadmin] schema cache

.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.

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,
@@ -185,6 +193,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-<key>' and as a gRPC metadata key '<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"))

go/vt/vtadmin/api.go

@@ -212,6 +212,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 {

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/vt/grpccommon"
 	"vitess.io/vitess/go/vt/topo"
@@ -99,6 +100,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{
 						{
@@ -108,6 +125,7 @@ func TestFindSchema(t *testing.T) {
 									Uid:  100,
 								},
 								Keyspace: "testkeyspace",
+								Shard:    "-",
 							},
 							State: vtadminpb.Tablet_SERVING,
 						},
@@ -530,6 +548,7 @@ func TestFindSchema(t *testing.T) {
 			}
 
 			api := NewAPI(clusters, Options{})
+			defer api.Close()
 
 			resp, err := api.FindSchema(ctx, tt.req)
 			if tt.shouldErr {
@@ -739,6 +758,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{
@@ -773,6 +794,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.
@@ -1378,7 +1402,6 @@ func TestGetSchema(t *testing.T) {
 						Name: "testtable",
 					},
 				},
-				TableSizes: map[string]*vtadminpb.Schema_TableSize{},
 			},
 			shouldErr: false,
 		},
@@ -1528,6 +1551,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 {
@@ -1536,6 +1560,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)
 			})
@@ -1654,6 +1683,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",
@@ -1690,6 +1721,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.
@@ -2204,6 +2238,7 @@ func TestGetSchemas(t *testing.T) {
 				}
 
 				api := NewAPI(clusters, Options{})
+				defer api.Close()
 
 				resp, err := api.GetSchemas(ctx, tt.req)
 				require.NoError(t, err)
@@ -2424,6 +2459,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,
@@ -2491,14 +2528,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)