Cache Bigtable Table Info (grafana#1651)

* cache table info for bigtable tables in tablemanager Signed-off-by: Jacob Lisi <jacob.t.lisi@gmail.com> * Add expiration to the bigtable table cache Signed-off-by: Jacob Lisi <jacob.t.lisi@gmail.com> * fix order of adding table to cache Signed-off-by: Jacob Lisi <jacob.t.lisi@gmail.com> * refactor based on review comments Signed-off-by: Jacob Lisi <jacob.t.lisi@gmail.com> * add ListTables comment Signed-off-by: Jacob Lisi <jacob.t.lisi@gmail.com>
cyriltovena · Oct 16, 2019 · c848882 · c848882
1 parent 4ce3d0f
commit c848882
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/gcp/bigtable_index_client.go b/gcp/bigtable_index_client.go
@@ -8,6 +8,7 @@ import (
 	"flag"
 	"fmt"
 	"strings"
+	"time"
 
 	"cloud.google.com/go/bigtable"
 	ot "github.com/opentracing/opentracing-go"
@@ -38,12 +39,17 @@ type Config struct {
 
 	ColumnKey      bool
 	DistributeKeys bool
+
+	TableCacheEnabled    bool
+	TableCacheExpiration time.Duration
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Project, "bigtable.project", "", "Bigtable project ID.")
 	f.StringVar(&cfg.Instance, "bigtable.instance", "", "Bigtable instance ID.")
+	f.BoolVar(&cfg.TableCacheEnabled, "bigtable.table-cache.enabled", true, "If enabled, once a tables info is fetched, it is cached.")
+	f.DurationVar(&cfg.TableCacheExpiration, "bigtable.table-cache.expiration", 30*time.Minute, "Duration to cache tables before checking again.")
 
 	cfg.GRPCClientConfig.RegisterFlags("bigtable", f)
 }

diff --git a/gcp/table_client.go b/gcp/table_client.go
@@ -2,6 +2,7 @@ package gcp
 
 import (
 	"context"
+	"time"
 
 	"google.golang.org/grpc/codes"
 
@@ -15,6 +16,9 @@ import (
 type tableClient struct {
 	cfg    Config
 	client *bigtable.AdminClient
+
+	tableInfo       map[string]*bigtable.TableInfo
+	tableExpiration time.Time
 }
 
 // NewTableClient returns a new TableClient.
@@ -27,25 +31,37 @@ func NewTableClient(ctx context.Context, cfg Config) (chunk.TableClient, error)
 	return &tableClient{
 		cfg:    cfg,
 		client: client,
+
+		tableInfo: map[string]*bigtable.TableInfo{},
 	}, nil
 }
 
+// ListTables lists all of the correctly specified cortex tables in bigtable
 func (c *tableClient) ListTables(ctx context.Context) ([]string, error) {
 	tables, err := c.client.Tables(ctx)
 	if err != nil {
 		return nil, errors.Wrap(err, "client.Tables")
 	}
 
-	// Check each table has the right column family.  If not, omit it.
+	if c.tableExpiration.Before(time.Now()) {
+		c.tableInfo = map[string]*bigtable.TableInfo{}
+		c.tableExpiration = time.Now().Add(c.cfg.TableCacheExpiration)
+	}
+
 	output := make([]string, 0, len(tables))
 	for _, table := range tables {
-		info, err := c.client.TableInfo(ctx, table)
-		if err != nil {
-			return nil, errors.Wrap(err, "client.TableInfo")
+		info, exists := c.tableInfo[table]
+		if !c.cfg.TableCacheEnabled || !exists {
+			info, err = c.client.TableInfo(ctx, table)
+			if err != nil {
+				return nil, errors.Wrap(err, "client.TableInfo")
+			}
 		}
 
+		// Check each table has the right column family.  If not, omit it.
 		if hasColumnFamily(info.FamilyInfos) {
 			output = append(output, table)
+			c.tableInfo[table] = info
 		}
 	}
 
@@ -86,6 +102,7 @@ func (c *tableClient) DeleteTable(ctx context.Context, name string) error {
 	if err := c.client.DeleteTable(ctx, name); err != nil {
 		return errors.Wrap(err, "client.DeleteTable")
 	}
+	delete(c.tableInfo, name)
 
 	return nil
 }