feat(cache)!: use cachestore to cache request in server

internal/app/server/triton.go

@@ -29,6 +29,9 @@ import (
 	empty "google.golang.org/protobuf/types/known/emptypb"
 )
 
+// TODO(hongalex): make this a configurable field for users.
+const maxRecordSizeToCache int = 10 * 1024 * 1024 // 10 MB
+
 type tritonServer struct {
 	cloud      string
 	blobStore  blob.BlobStore
@@ -37,7 +40,7 @@ type tritonServer struct {
 }
 
 // newTritonServer creates a new instance of the triton server.
-func newTritonServer(ctx context.Context, cloud, project, bucket, cacheAddr string) (tritonpb.TritonServer, error) {
+func newTritonServer(ctx context.Context, cloud, project, bucket, cacheAddr string) (*tritonServer, error) {
 	switch cloud {
 	case "gcp":
 		log.Infoln("Instantiating Triton server on GCP")
@@ -91,8 +94,22 @@ func (s *tritonServer) CreateRecord(ctx context.Context, req *tritonpb.CreateRec
 			req.GetStoreKey(), req.Record.GetKey(), err)
 		return nil, status.Convert(err).Err()
 	}
-	log.Infof("Created record: %+v", record)
-	return newRecord.ToProto(), nil
+
+	// Update cache store.
+	k := cache.FormatKey(req.GetStoreKey(), req.Record.GetKey())
+	rp := newRecord.ToProto()
+	by, err := cache.EncodeRecord(rp)
+	if err != nil {
+		// Cache fails should be logged but not return error.
+		log.Errorf("failed to encode record for cache for key (%s): %v", k, err)
+	} else {
+		if len(by) < maxRecordSizeToCache {
+			if err := s.cacheStore.Set(ctx, k, by); err != nil {
+				log.Errorf("failed to update cache for key (%s): %v", k, err)
+			}
+		}
+	}
+	return rp, nil
 }
 
 func (s *tritonServer) DeleteRecord(ctx context.Context, req *tritonpb.DeleteRecordRequest) (*empty.Empty, error) {
@@ -104,6 +121,13 @@ func (s *tritonServer) DeleteRecord(ctx context.Context, req *tritonpb.DeleteRec
 	}
 	log.Infof("Deleted record: store (%s), record (%s)",
 		req.GetStoreKey(), req.GetKey())
+
+	// Purge record from cache store.
+	k := cache.FormatKey(req.GetStoreKey(), req.GetKey())
+	if err := s.cacheStore.Delete(ctx, k); err != nil {
+		log.Errorf("failed to purge cache for key (%s): %v", k, err)
+	}
+
 	return new(empty.Empty), nil
 }
 
@@ -140,14 +164,42 @@ func (s *tritonServer) DeleteStore(ctx context.Context, req *tritonpb.DeleteStor
 }
 
 func (s *tritonServer) GetRecord(ctx context.Context, req *tritonpb.GetRecordRequest) (*tritonpb.Record, error) {
+	k := cache.FormatKey(req.GetStoreKey(), req.GetKey())
+	r, err := s.cacheStore.Get(ctx, k)
+
+	// Cache hit, use value from cache store.
+	if err == nil {
+		re, err := cache.DecodeRecord(r)
+		if err != nil {
+			return nil, err
+		}
+		log.Infof("cache hit: %+v", re)
+		return re, nil
+	}
+
 	record, err := s.metaDB.GetRecord(ctx, req.GetStoreKey(), req.GetKey())
 	if err != nil {
 		log.Warnf("GetRecord failed for store (%s), record (%s): %v",
 			req.GetStoreKey(), req.GetKey(), err)
 		return nil, status.Convert(err).Err()
 	}
-	log.Infof("Got record: %+v", record)
-	return record.ToProto(), nil
+	log.Infof("Got record %+v", record)
+
+	// Update cache store.
+	rp := record.ToProto()
+	by, err := cache.EncodeRecord(rp)
+	if err != nil {
+		// Cache fails should be logged but not return error.
+		log.Errorf("failed to encode record for cache for key (%s): %v", k, err)
+	} else {
+		if len(by) < maxRecordSizeToCache {
+			if err := s.cacheStore.Set(ctx, k, by); err != nil {
+				log.Errorf("failed to update cache for key (%s): %v", k, err)
+			}
+		}
+	}
+
+	return rp, nil
 }
 
 func (s *tritonServer) UpdateRecord(ctx context.Context, req *tritonpb.UpdateRecordRequest) (*tritonpb.Record, error) {
@@ -158,7 +210,23 @@ func (s *tritonServer) UpdateRecord(ctx context.Context, req *tritonpb.UpdateRec
 			req.GetStoreKey(), req.GetRecord().GetKey(), err)
 		return nil, status.Convert(err).Err()
 	}
-	return newRecord.ToProto(), nil
+
+	// Update cache store.
+	k := cache.FormatKey(req.GetStoreKey(), req.GetRecord().GetKey())
+	rp := newRecord.ToProto()
+	by, err := cache.EncodeRecord(rp)
+	if err != nil {
+		// Cache fails should be logged but not return error.
+		log.Errorf("failed to encode record for cache for key (%s): %v", k, err)
+	} else {
+		if len(by) < maxRecordSizeToCache {
+			if err := s.cacheStore.Set(ctx, k, by); err != nil {
+				log.Errorf("failed to update cache for key (%s): %v", k, err)
+			}
+		}
+	}
+
+	return rp, nil
 }
 
 func (s *tritonServer) QueryRecords(ctx context.Context, req *tritonpb.QueryRecordsRequest) (*tritonpb.QueryRecordsResponse, error) {

internal/pkg/cache/cache.go

@@ -14,11 +14,61 @@
 
 package cache
 
-import "context"
+import (
+	"bytes"
+	"context"
+	"encoding/gob"
+	"fmt"
+	"sync"
+
+	"github.com/golang/protobuf/ptypes/timestamp"
+	tritonpb "github.com/googleforgames/triton/api"
+)
 
 type Cache interface {
-	Set(ctx context.Context, key string, value string) error
-	Get(ctx context.Context, key string) (string, error)
+	Set(ctx context.Context, key string, value []byte) error
+	Get(ctx context.Context, key string) ([]byte, error)
 	Delete(ctx context.Context, key string) error
 	ListKeys(ctx context.Context) ([]string, error)
+	FlushAll(ctx context.Context) error
+}
+
+var once sync.Once
+
+// registerProperties is called once and used to register new types
+// for gob encoding/decoding.
+func registerProperties() {
+	gob.Register(&tritonpb.Property_BooleanValue{})
+	gob.Register(&tritonpb.Property_IntegerValue{})
+	gob.Register(&tritonpb.Property_StringValue{})
+	gob.Register(&timestamp.Timestamp{})
+}
+
+// FormatKey concatenates store and record keys separated by a backslash.
+func FormatKey(storeKey, recordKey string) string {
+	return fmt.Sprintf("%s/%s", storeKey, recordKey)
+}
+
+// EncodeRecord serializes a tritonpb Record with gob.
+func EncodeRecord(r *tritonpb.Record) ([]byte, error) {
+	once.Do(registerProperties)
+	b := bytes.Buffer{}
+	e := gob.NewEncoder(&b)
+	if err := e.Encode(r); err != nil {
+		return nil, err
+	}
+	return b.Bytes(), nil
+}
+
+// EncodeRecord deserializes a tritonpb Record with gob.
+func DecodeRecord(by []byte) (*tritonpb.Record, error) {
+	once.Do(registerProperties)
+	r := &tritonpb.Record{}
+	b := bytes.Buffer{}
+	b.Write(by)
+	d := gob.NewDecoder(&b)
+	if err := d.Decode(&r); err != nil {
+		return nil, err
+	}
+	return r, nil
 }

internal/pkg/cache/cache_test.go

@@ -0,0 +1,122 @@
+// Copyright 2020 Google LLC
+//
+// 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 (
+	"testing"
+	"time"
+
+	"github.com/golang/protobuf/ptypes/timestamp"
+	pb "github.com/googleforgames/triton/api"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	// The threshold of comparing times.
+	// Since the server and client run on the same host for these tests,
+	// 1 second should be enough.
+	timestampDelta = 1 * time.Second
+)
+
+func TestCache_SerializeRecord(t *testing.T) {
+	rr := []*pb.Record{
+		{
+			CreatedAt: &timestamp.Timestamp{
+				Seconds: 100,
+			},
+			UpdatedAt: &timestamp.Timestamp{
+				Seconds: 110,
+			},
+		},
+		{
+			Key: "some-key",
+			Properties: map[string]*pb.Property{
+				"prop1": {
+					Type: pb.Property_BOOLEAN,
+					Value: &pb.Property_BooleanValue{
+						BooleanValue: false,
+					},
+				},
+				"prop2": {
+					Type: pb.Property_INTEGER,
+					Value: &pb.Property_IntegerValue{
+						IntegerValue: 200,
+					},
+				},
+				"prop3": {
+					Type: pb.Property_STRING,
+					Value: &pb.Property_StringValue{
+						StringValue: "triton",
+					},
+				},
+			},
+			CreatedAt: &timestamp.Timestamp{
+				Seconds: 100,
+			},
+			UpdatedAt: &timestamp.Timestamp{
+				Seconds: 110,
+			},
+		},
+		{
+			Key:      "some-key",
+			Blob:     []byte("some-bytes"),
+			BlobSize: 64,
+			OwnerId:  "new-owner",
+			Tags:     []string{"tag1", "tag2"},
+			CreatedAt: &timestamp.Timestamp{
+				Seconds: 100,
+			},
+			UpdatedAt: &timestamp.Timestamp{
+				Seconds: 110,
+			},
+		},
+	}
+
+	for _, r := range rr {
+		e, err := EncodeRecord(r)
+		assert.NoError(t, err)
+		d, err := DecodeRecord(e)
+		assert.NoError(t, err)
+		assertEqualRecord(t, r, d)
+	}
+}
+
+func assertEqualRecord(t *testing.T, expected, actual *pb.Record) {
+	if expected == nil {
+		assert.Nil(t, actual)
+		return
+	}
+	if assert.NotNil(t, actual) {
+		assert.Equal(t, expected.Key, actual.Key)
+		assert.Equal(t, expected.Blob, actual.Blob)
+		assert.Equal(t, expected.BlobSize, actual.BlobSize)
+		assert.Equal(t, expected.Tags, actual.Tags)
+		assert.Equal(t, expected.OwnerId, actual.OwnerId)
+		assert.Equal(t, len(expected.Properties), len(actual.Properties))
+		for k, v := range expected.Properties {
+			if assert.Contains(t, actual.Properties, k) {
+				av := actual.Properties[k]
+				assert.Equal(t, v.Type, av.Type)
+				assert.Equal(t, v.Value, av.Value)
+			}
+		}
+		assert.NotNil(t, actual.GetCreatedAt())
+		assert.WithinDuration(t, expected.GetUpdatedAt().AsTime(),
+			actual.GetUpdatedAt().AsTime(), timestampDelta)
+		assert.NotNil(t, actual.GetUpdatedAt())
+		assert.WithinDuration(t, expected.GetUpdatedAt().AsTime(),
+			actual.GetUpdatedAt().AsTime(), timestampDelta)
+	}
+}

internal/pkg/cache/redis.go

@@ -37,7 +37,7 @@ func NewRedis(address string, opts ...redis.DialOption) *Redis {
 	}
 }
 
-func (r *Redis) Set(ctx context.Context, key, value string) error {
+func (r *Redis) Set(ctx context.Context, key string, value []byte) error {
 	conn := r.redisPool.Get()
 	defer conn.Close()
 
@@ -48,13 +48,13 @@ func (r *Redis) Set(ctx context.Context, key, value string) error {
 	return nil
 }
 
-func (r *Redis) Get(ctx context.Context, key string) (string, error) {
+func (r *Redis) Get(ctx context.Context, key string) ([]byte, error) {
 	conn := r.redisPool.Get()
 	defer conn.Close()
 
-	val, err := redis.String(conn.Do("GET", key))
+	val, err := redis.Bytes(conn.Do("GET", key))
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	return val, nil
 }
@@ -70,6 +70,17 @@ func (r *Redis) Delete(ctx context.Context, key string) error {
 	return nil
 }
 
+func (r *Redis) FlushAll(ctx context.Context) error {
+	conn := r.redisPool.Get()
+	defer conn.Close()
+
+	_, err := conn.Do("FLUSHALL")
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (r *Redis) ListKeys(ctx context.Context) ([]string, error) {
 	conn := r.redisPool.Get()
 	defer conn.Close()