Add CloseSession RPCjc-evict-batch-cache

8 files changed, 130 insertions, 11 deletions

diff --git a/internal/git/catfile/batch_cache.go b/internal/git/catfile/batch_cache.go
index 10ea489a1..daa006f72 100644
--- a/internal/git/catfile/batch_cache.go
+++ b/internal/git/catfile/batch_cache.go
@@ -157,6 +157,24 @@ func (bc *batchCache) EvictAll() {
 	}
 }
 
+func EvictCacheBySessionID(sessionID string) {
+	cache.EvictBySessionID(sessionID)
+}
+
+func (bc *batchCache) EvictBySessionID(sessionID string) {
+	bc.Lock()
+	defer bc.Unlock()
+
+	for i := 0; i < len(bc.entries); i++ {
+		ent := bc.entries[i]
+		if ent.key.sessionID == sessionID {
+			bc.delete(i, true)
+			i--
+		}
+	}
+	catfileCacheMembers.Set(float64(bc.len()))
+}
+
 // ExpireAll is used to expire all of the batches in the cache
 func ExpireAll() {
 	cache.EvictAll()
diff --git a/internal/git/catfile/batch_cache_test.go b/internal/git/catfile/batch_cache_test.go
index 0f885d314..c437066e4 100644
--- a/internal/git/catfile/batch_cache_test.go
+++ b/internal/git/catfile/batch_cache_test.go
@@ -153,6 +153,41 @@ func TestAutoExpiry(t *testing.T) {
 	bc.Unlock()
 }
 
+func TestExpireBySessionID(t *testing.T) {
+
+	ttl := 5 * time.Millisecond
+	bc := newCacheWithRefresh(ttl, 10, 1*time.Millisecond)
+
+	keysToGetEvicted := []key{
+		key{sessionID: "a", repoRelPath: "path_a"},
+		key{sessionID: "a", repoRelPath: "path_b"},
+	}
+
+	keysToRemain := []key{
+		key{sessionID: "b", repoRelPath: "path_b"},
+		key{sessionID: "c", repoRelPath: "path_b"},
+		key{sessionID: "d", repoRelPath: "path_b"},
+		key{sessionID: "e", repoRelPath: "path_b"},
+		key{sessionID: "f", repoRelPath: "path_b"},
+	}
+
+	for _, key := range append(keysToGetEvicted, keysToRemain...) {
+		bc.Add(key, testValue())
+	}
+
+	bc.EvictBySessionID("a")
+
+	for _, k := range keysToGetEvicted {
+		_, ok := bc.Checkout(k)
+		require.False(t, ok)
+	}
+
+	for _, k := range keysToRemain {
+		_, ok := bc.Checkout(k)
+		require.True(t, ok)
+	}
+}
+
 func requireCacheValid(t *testing.T, bc *batchCache) {
 	bc.Lock()
 	defer bc.Unlock()
diff --git a/internal/service/cleanup/apply_bfg_object_map_stream_test.go b/internal/service/cleanup/apply_bfg_object_map_stream_test.go
index e4c916330..d1d9baefc 100644
--- a/internal/service/cleanup/apply_bfg_object_map_stream_test.go
+++ b/internal/service/cleanup/apply_bfg_object_map_stream_test.go
@@ -1,4 +1,4 @@
-package cleanup
+package cleanup_test
 
 import (
 	"context"
diff --git a/internal/service/cleanup/apply_bfg_object_map_test.go b/internal/service/cleanup/apply_bfg_object_map_test.go
index 9d1e087ff..faede8d5b 100644
--- a/internal/service/cleanup/apply_bfg_object_map_test.go
+++ b/internal/service/cleanup/apply_bfg_object_map_test.go
@@ -1,4 +1,4 @@
-package cleanup
+package cleanup_test
 
 import (
 	"context"
diff --git a/internal/service/cleanup/close_session.go b/internal/service/cleanup/close_session.go
new file mode 100644
index 000000000..784068b12
--- /dev/null
+++ b/internal/service/cleanup/close_session.go
@@ -0,0 +1,29 @@
+package cleanup
+
+import (
+	"context"
+	"errors"
+
+	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
+	"gitlab.com/gitlab-org/gitaly/internal/helper"
+
+	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
+)
+
+func (s *server) CloseSession(ctx context.Context, in *gitalypb.CloseSessionRequest) (*gitalypb.CloseSessionResponse, error) {
+	if err := validateCloseSessionRequest(in); err != nil {
+		return nil, helper.ErrInvalidArgument(err)
+	}
+
+	catfile.EvictCacheBySessionID(in.GetSessionId())
+
+	return &gitalypb.CloseSessionResponse{}, nil
+}
+
+func validateCloseSessionRequest(in *gitalypb.CloseSessionRequest) error {
+	if in.GetSessionId() == "" {
+		return errors.New("session_id is required")
+	}
+
+	return nil
+}
diff --git a/internal/service/cleanup/close_session_test.go b/internal/service/cleanup/close_session_test.go
new file mode 100644
index 000000000..2b1921a7b
--- /dev/null
+++ b/internal/service/cleanup/close_session_test.go
@@ -0,0 +1,43 @@
+package cleanup_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
+	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestCloseSession(t *testing.T) {
+	server, serverSocketPath := runCleanupServiceServer(t)
+	defer server.Stop()
+
+	client, conn := newCleanupServiceClient(t, serverSocketPath)
+	defer conn.Close()
+
+	testRepo, _, cleanupFn := testhelper.NewTestRepo(t)
+	defer cleanupFn()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	md := metadata.New(map[string]string{
+		featureflag.HeaderKey(catfile.CacheFeatureFlagKey): "true",
+		"gitaly-session-id": t.Name(),
+	})
+
+	ctx = metadata.NewOutgoingContext(ctx, md)
+
+	catfile.New(ctx, testRepo)
+
+	ctx, cancel = testhelper.Context()
+	defer cancel()
+	_, err := client.CloseSession(ctx, &gitalypb.CloseSessionRequest{
+		SessionId: t.Name(),
+	})
+
+	require.NoError(t, err)
+}
diff --git a/internal/service/cleanup/server.go b/internal/service/cleanup/server.go
index b1bd4e0f5..c5bcb6d8f 100644
--- a/internal/service/cleanup/server.go
+++ b/internal/service/cleanup/server.go
@@ -1,10 +1,7 @@
 package cleanup
 
 import (
-	"context"
-
 	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
-	"gitlab.com/gitlab-org/gitaly/internal/helper"
 )
 
 type server struct{}
@@ -13,7 +10,3 @@ type server struct{}
 func NewServer() gitalypb.CleanupServiceServer {
 	return &server{}
 }
-
-func (s *server) CloseSession(context.Context, *gitalypb.CloseSessionRequest) (*gitalypb.CloseSessionResponse, error) {
-	return nil, helper.Unimplemented
-}
diff --git a/internal/service/cleanup/testhelper_test.go b/internal/service/cleanup/testhelper_test.go
index a43a304f1..46855a12a 100644
--- a/internal/service/cleanup/testhelper_test.go
+++ b/internal/service/cleanup/testhelper_test.go
@@ -1,4 +1,4 @@
-package cleanup
+package cleanup_test
 
 import (
 	"net"
@@ -8,6 +8,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
 
+	"gitlab.com/gitlab-org/gitaly/internal/service/cleanup"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 )
 
@@ -20,7 +21,7 @@ func runCleanupServiceServer(t *testing.T) (*grpc.Server, string) {
 		t.Fatal(err)
 	}
 
-	gitalypb.RegisterCleanupServiceServer(grpcServer, NewServer())
+	gitalypb.RegisterCleanupServiceServer(grpcServer, cleanup.NewServer())
 	reflection.Register(grpcServer)
 
 	go grpcServer.Serve(listener)