Implement cache for VV in MongoDB client (#1231)

Introduced an LRU cache in the MongoDB client to store version vectors
per document, reducing redundant DB queries. The cache is initialized
on client creation and used in UpdateAndFindMinSyncedVersionVector.

Author: hackerwins

server/backend/database/database.go

@@ -58,6 +58,9 @@ var (
 
 	// ErrProjectNameAlreadyExists is returned when the project name already exists.
 	ErrProjectNameAlreadyExists = errors.New("project name already exists")
+
+	// ErrVersionVectorNotFound is returned when the version vector could not be found.
+	ErrVersionVectorNotFound = errors.New("version vector not found")
 )
 
 // Database represents database which reads or saves Yorkie data.

server/backend/database/mongo/client.go

@@ -25,6 +25,7 @@ import (
 	"strings"
 	gotime "time"
 
+	lru "github.com/hashicorp/golang-lru/v2"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -33,6 +34,7 @@ import (
 
 	"github.com/yorkie-team/yorkie/api/converter"
 	"github.com/yorkie-team/yorkie/api/types"
+	"github.com/yorkie-team/yorkie/pkg/cmap"
 	"github.com/yorkie-team/yorkie/pkg/document"
 	"github.com/yorkie-team/yorkie/pkg/document/change"
 	"github.com/yorkie-team/yorkie/pkg/document/key"
@@ -48,8 +50,9 @@ const (
 
 // Client is a client that connects to Mongo DB and reads or saves Yorkie data.
 type Client struct {
-	config *Config
-	client *mongo.Client
+	config  *Config
+	client  *mongo.Client
+	vvCache *lru.Cache[types.DocRefKey, *cmap.Map[types.ID, time.VersionVector]]
 }
 
 // Dial creates an instance of Client and dials the given MongoDB.
@@ -79,11 +82,17 @@ func Dial(conf *Config) (*Client, error) {
 		return nil, err
 	}
 
+	cache, err := lru.New[types.DocRefKey, *cmap.Map[types.ID, time.VersionVector]](1000)
+	if err != nil {
+		return nil, fmt.Errorf("initialize VV cache: %w", err)
+	}
+
 	logging.DefaultLogger().Infof("MongoDB connected, URI: %s, DB: %s", conf.ConnectionURI, conf.YorkieDatabase)
 
 	return &Client{
-		config: conf,
-		client: client,
+		config:  conf,
+		client:  client,
+		vvCache: cache,
 	}, nil
 }
 
@@ -93,6 +102,8 @@ func (c *Client) Close() error {
 		return fmt.Errorf("close mongo client: %w", err)
 	}
 
+	c.vvCache.Purge()
+
 	return nil
 }
 
@@ -1314,54 +1325,79 @@ func (c *Client) UpdateAndFindMinSyncedTicket(
 	), nil
 }
 
-// UpdateAndFindMinSyncedVersionVector returns min synced version vector
+// UpdateAndFindMinSyncedVersionVector returns min synced version vector.
 func (c *Client) UpdateAndFindMinSyncedVersionVector(
 	ctx context.Context,
 	clientInfo *database.ClientInfo,
 	docRefKey types.DocRefKey,
-	versionVector time.VersionVector,
+	vector time.VersionVector,
 ) (time.VersionVector, error) {
-	// TODO(JOOHOJANG): We have to consider removing detached client's lamport
-	// from min version vector.
-	var versionVectorInfos []database.VersionVectorInfo
-
-	// 01. Find all version vectors of the given document from DB.
-	cursor, err := c.collection(ColVersionVectors).Find(ctx, bson.M{
-		"project_id": docRefKey.ProjectID,
-		"doc_id":     docRefKey.DocID,
-	})
-	if err != nil {
-		return nil, fmt.Errorf("find all version vectors: %w", err)
+	// 01. Update synced version vector of the given client and document.
+	if err := c.UpdateVersionVector(ctx, clientInfo, docRefKey, vector); err != nil {
+		return nil, err
 	}
 
-	if err := cursor.All(ctx, &versionVectorInfos); err != nil {
-		return nil, fmt.Errorf("decode version vectors: %w", err)
+	// 02. Update current client's version vector. If the client is detached, remove it.
+	// This is only for the current client and does not affect the version vector of other clients.
+	if vvMap, ok := c.vvCache.Get(docRefKey); ok {
+		attached, err := clientInfo.IsAttached(docRefKey.DocID)
+		if err != nil {
+			return nil, err
+		}
+
+		if attached {
+			vvMap.Upsert(clientInfo.ID, func(value time.VersionVector, exists bool) time.VersionVector {
+				return vector
+			})
+		} else {
+			// NOTE(hackerwins): Considering removing the detached client's lamport
+			// from the other clients' version vectors. For now, we just ignore it.
+			vvMap.Delete(clientInfo.ID, func(value time.VersionVector, exists bool) bool {
+				return exists
+			})
+		}
 	}
 
-	// 02. Compute min version vector.
-	minVersionVector := versionVector.DeepCopy()
-	for i, vvi := range versionVectorInfos {
-		if vvi.ClientID == clientInfo.ID {
-			continue
+	// 03. Calculate the minimum version vector of the given document.
+	if !c.vvCache.Contains(docRefKey) {
+		var infos []database.VersionVectorInfo
+		cursor, err := c.collection(ColVersionVectors).Find(ctx, bson.M{
+			"project_id": docRefKey.ProjectID,
+			"doc_id":     docRefKey.DocID,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("find all version vectors: %w", err)
+		}
+		if err := cursor.All(ctx, &infos); err != nil {
+			return nil, fmt.Errorf("decode version vectors: %w", err)
+		}
+
+		infoMap := cmap.New[types.ID, time.VersionVector]()
+		for i := range infos {
+			infoMap.Set(infos[i].ClientID, infos[i].VersionVector)
 		}
-		minVersionVector.Min(&versionVectorInfos[i].VersionVector)
+
+		c.vvCache.Add(docRefKey, infoMap)
+	}
+	vvMap, ok := c.vvCache.Get(docRefKey)
+	if !ok {
+		return nil, fmt.Errorf("version vectors from cache: %w", database.ErrVersionVectorNotFound)
 	}
 
-	// 03. Update current client's version vector. If the client is detached, remove it.
-	// This is only for the current client and does not affect the version vector of other clients.
-	if err = c.UpdateVersionVector(ctx, clientInfo, docRefKey, versionVector); err != nil {
-		return nil, err
+	minVector := vector.DeepCopy()
+	for _, vv := range vvMap.Values() {
+		minVector.Min(&vv)
 	}
 
-	return minVersionVector, nil
+	return minVector, nil
 }
 
 // UpdateVersionVector updates the given version vector of the given client
 func (c *Client) UpdateVersionVector(
 	ctx context.Context,
 	clientInfo *database.ClientInfo,
 	docRefKey types.DocRefKey,
-	versionVector time.VersionVector,
+	vector time.VersionVector,
 ) error {
 	isAttached, err := clientInfo.IsAttached(docRefKey.DocID)
 	if err != nil {
@@ -1385,7 +1421,7 @@ func (c *Client) UpdateVersionVector(
 		"client_id":  clientInfo.ID,
 	}, bson.M{
 		"$set": bson.M{
-			"version_vector": versionVector,
+			"version_vector": vector,
 		},
 	}, options.Update().SetUpsert(true))
 	if err != nil {