[ENH]: Modified AttachFunction to do 2PC on a new is_ready column

Author: tanujnay112

chromadb/test/distributed/test_task_api.py

@@ -18,7 +18,6 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
-	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // validateAttachedFunctionMatchesRequest validates that an existing attached function's parameters match the request parameters.
@@ -88,10 +87,8 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 	}
 
 	var attachedFunctionID uuid.UUID = uuid.New()
-	var nextRun time.Time
-	var skipPhase2 bool // Flag to skip Phase 2 if task is already fully created
 
-	// ===== Phase 1: Create attached function (if needed) =====
+	// ===== Step 1: Create attached function with is_ready = false =====
 	err := s.catalog.txImpl.Transaction(ctx, func(txCtx context.Context) error {
 		// Double-check attached function doesn't exist (race condition protection)
 		concurrentAttachedFunction, err := s.catalog.metaDomain.AttachedFunctionDb(txCtx).GetByName(req.InputCollectionId, req.Name)
@@ -109,7 +106,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 				return err
 			}
 
-			// Validation passed, reuse the concurrent attached function's data
+			// Validation passed, reuse the concurrent attached function ID (idempotent)
 			attachedFunctionID = concurrentAttachedFunction.ID
 			// Already created, skip Phase 2
 			skipPhase2 = true
@@ -137,7 +134,6 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			log.Error("AttachFunction: function not found", zap.String("function_name", req.FunctionName))
 			return common.ErrFunctionNotFound
 		}
-
 		// Check if input collection exists
 		collections, err := s.catalog.metaDomain.CollectionDb(txCtx).GetCollections([]string{req.InputCollectionId}, nil, req.TenantId, req.Database, nil, nil, false)
 		if err != nil {
@@ -191,6 +187,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			CreatedAt:               now,
 			UpdatedAt:               now,
 			OldestWrittenNonce:      nil,
+			IsReady:                 false, // **KEY: Set to false for 3-phase create**
 		}
 
 		err = s.catalog.metaDomain.AttachedFunctionDb(txCtx).Insert(attachedFunction)
@@ -199,7 +196,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			return err
 		}
 
-		log.Debug("AttachFunction: Phase 1: attached function created",
+		log.Debug("AttachFunction: attached function created with is_ready=false",
 			zap.String("attached_function_id", attachedFunctionID.String()),
 			zap.String("name", req.Name))
 		return nil
@@ -209,45 +206,6 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 		return nil, err
 	}
 
-	// If function is already created, return immediately (idempotency)
-	if skipPhase2 {
-		log.Info("AttachFunction: function already created, skipping Phase 2",
-			zap.String("attached_function_id", attachedFunctionID.String()))
-		return &coordinatorpb.AttachFunctionResponse{
-			Id: attachedFunctionID.String(),
-		}, nil
-	}
-
-	// ===== Phase 2: Push initial schedule =====
-	log.Debug("AttachFunction: Phase 2: pushing initial schedule",
-		zap.String("attached_function_id", attachedFunctionID.String()))
-	// Push initial schedule to heap service if enabled
-	if s.heapClient == nil {
-		return nil, common.ErrHeapServiceNotEnabled
-	}
-
-	// Create schedule for the attached function
-	schedule := &coordinatorpb.Schedule{
-		Triggerable: &coordinatorpb.Triggerable{
-			PartitioningUuid: req.InputCollectionId,
-			SchedulingUuid:   attachedFunctionID.String(),
-		},
-		NextScheduled: timestamppb.New(nextRun),
-	}
-
-	err = s.heapClient.Push(ctx, req.InputCollectionId, []*coordinatorpb.Schedule{schedule})
-	if err != nil {
-		log.Error("AttachFunction: Phase 2: failed to push schedule to heap service",
-			zap.Error(err),
-			zap.String("attached_function_id", attachedFunctionID.String()),
-			zap.String("collection_id", req.InputCollectionId))
-		return nil, err
-	}
-
-	log.Debug("AttachFunction: Phase 2: pushed schedule to heap service",
-		zap.String("attached_function_id", attachedFunctionID.String()),
-		zap.String("collection_id", req.InputCollectionId))
-
 	return &coordinatorpb.AttachFunctionResponse{
 		Id: attachedFunctionID.String(),
 	}, nil
@@ -287,6 +245,7 @@ func attachedFunctionToProto(attachedFunction *dbmodel.AttachedFunction, functio
 		DatabaseId:              attachedFunction.DatabaseID,
 		CreatedAt:               uint64(attachedFunction.CreatedAt.UnixMicro()),
 		UpdatedAt:               uint64(attachedFunction.UpdatedAt.UnixMicro()),
+		IsReady:                 attachedFunction.IsReady,
 	}
 	if attachedFunction.OutputCollectionID != nil {
 		attachedFunctionProto.OutputCollectionId = attachedFunction.OutputCollectionID
@@ -563,7 +522,6 @@ func (s *Coordinator) DetachFunction(ctx context.Context, req *coordinatorpb.Det
 	// First get the attached function to check if we need to delete the output collection
 	attachedFunction, err := s.catalog.metaDomain.AttachedFunctionDb(ctx).GetByID(attachedFunctionID)
 	if err != nil {
-		// If attached function is not ready (lowest_live_nonce == NULL), treat it as not found
 		if errors.Is(err, common.ErrAttachedFunctionNotReady) {
 			log.Error("DetachFunction: attached function not ready (not initialized)")
 			return nil, status.Error(codes.NotFound, "attached function not found")
@@ -654,6 +612,134 @@ func (s *Coordinator) GetFunctions(ctx context.Context, req *coordinatorpb.GetFu
 	}, nil
 }
 
+// TODO(tanujnay112): Remove this
+func (s *Coordinator) FinishAttachedFunction(ctx context.Context, req *coordinatorpb.FinishAttachedFunctionRequest) (*coordinatorpb.FinishAttachedFunctionResponse, error) {
+	attachedFunctionID, err := uuid.Parse(req.Id)
+	if err != nil {
+		log.Error("FinishAttachedFunction: invalid attached_function_id", zap.Error(err))
+		return nil, err
+	}
+
+	err = s.catalog.metaDomain.AttachedFunctionDb(ctx).Finish(attachedFunctionID)
+	if err != nil {
+		log.Error("FinishAttachedFunction: failed to finish attached function", zap.Error(err))
+		return nil, err
+	}
+
+	return &coordinatorpb.FinishAttachedFunctionResponse{}, nil
+}
+
+// FinishCreateAttachedFunction creates the output collection and sets is_ready to true in a single transaction
+func (s *Coordinator) FinishCreateAttachedFunction(ctx context.Context, req *coordinatorpb.FinishCreateAttachedFunctionRequest) (*coordinatorpb.FinishCreateAttachedFunctionResponse, error) {
+	attachedFunctionID, err := uuid.Parse(req.Id)
+	if err != nil {
+		log.Error("FinishCreateAttachedFunction: invalid attached_function_id", zap.Error(err))
+		return nil, status.Errorf(codes.InvalidArgument, "invalid attached_function_id: %v", err)
+	}
+
+	// Execute all operations in a transaction for atomicity
+	err = s.catalog.txImpl.Transaction(ctx, func(txCtx context.Context) error {
+		// 1. Get the attached function to retrieve metadata
+		attachedFunction, err := s.catalog.metaDomain.AttachedFunctionDb(txCtx).GetByID(attachedFunctionID)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to get attached function", zap.Error(err))
+			return err
+		}
+		if attachedFunction == nil {
+			log.Error("FinishCreateAttachedFunction: attached function not found")
+			return status.Errorf(codes.NotFound, "attached function not found")
+		}
+
+		// 2. Check if output collection already exists (idempotency)
+		if attachedFunction.OutputCollectionID != nil && *attachedFunction.OutputCollectionID != "" {
+			log.Info("FinishCreateAttachedFunction: output collection already exists, skipping creation",
+				zap.String("existing_collection_id", *attachedFunction.OutputCollectionID))
+			// Still set is_ready in case it wasn't set before
+			return s.catalog.metaDomain.AttachedFunctionDb(txCtx).SetReady(attachedFunctionID)
+		}
+
+		// 3. Look up database by ID to get its name
+		database, err := s.catalog.metaDomain.DatabaseDb(txCtx).GetByID(attachedFunction.DatabaseID)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to get database", zap.Error(err))
+			return err
+		}
+		if database == nil {
+			log.Error("FinishCreateAttachedFunction: database not found", zap.String("database_id", attachedFunction.DatabaseID), zap.String("tenant_id", attachedFunction.TenantID))
+			return common.ErrDatabaseNotFound
+		}
+
+		// 4. Generate new collection UUID
+		collectionID := types.NewUniqueID()
+
+		// 5. Create the output collection with segments
+		dimension := int32(1) // Default dimension for attached function output collections
+		collection := &model.CreateCollection{
+			ID:                   collectionID,
+			Name:                 attachedFunction.OutputCollectionName,
+			ConfigurationJsonStr: "{}", // Empty JSON object for default config
+			TenantID:             attachedFunction.TenantID,
+			DatabaseName:         database.Name,
+			Dimension:            &dimension,
+			Metadata:             nil,
+		}
+
+		// Create segments for the collection (distributed setup)
+		segments := []*model.Segment{
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/vector/hnsw-distributed",
+				Scope:        "VECTOR",
+				CollectionID: collectionID,
+			},
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/metadata/blockfile",
+				Scope:        "METADATA",
+				CollectionID: collectionID,
+			},
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/record/blockfile",
+				Scope:        "RECORD",
+				CollectionID: collectionID,
+			},
+		}
+
+		_, _, err = s.catalog.CreateCollectionAndSegments(txCtx, collection, segments, 0)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to create collection", zap.Error(err))
+			return err
+		}
+
+		// 6. Update attached function with output_collection_id
+		collectionIDStr := collectionID.String()
+		err = s.catalog.metaDomain.AttachedFunctionDb(txCtx).UpdateOutputCollectionID(attachedFunctionID, &collectionIDStr)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to update output collection ID", zap.Error(err))
+			return err
+		}
+
+		// 7. Set is_ready to true
+		err = s.catalog.metaDomain.AttachedFunctionDb(txCtx).SetReady(attachedFunctionID)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to set ready", zap.Error(err))
+			return err
+		}
+
+		log.Info("FinishCreateAttachedFunction: successfully created output collection and set is_ready=true",
+			zap.String("attached_function_id", attachedFunctionID.String()),
+			zap.String("output_collection_id", collectionID.String()))
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &coordinatorpb.FinishCreateAttachedFunctionResponse{}, nil
+}
+
 // CleanupExpiredPartialAttachedFunctions finds and soft deletes attached functions that were partially created
 // (output_collection_id IS NULL) and are older than the specified max age.
 // This is used to clean up attached functions that got stuck during creation.

go/pkg/sysdb/coordinator/task.go

@@ -48,6 +48,7 @@ func (s *attachedFunctionDb) GetByName(inputCollectionID string, name string) (*
 		Where("input_collection_id = ?", inputCollectionID).
 		Where("name = ?", name).
 		Where("is_deleted = ?", false).
+		Where("is_ready = ?", true).
 		First(&attachedFunction).Error
 
 	if err != nil {
@@ -66,6 +67,7 @@ func (s *attachedFunctionDb) GetByID(id uuid.UUID) (*dbmodel.AttachedFunction, e
 	err := s.db.
 		Where("id = ?", id).
 		Where("is_deleted = ?", false).
+		Where("is_ready = ?", true).
 		First(&attachedFunction).Error
 
 	if err != nil {
@@ -84,6 +86,7 @@ func (s *attachedFunctionDb) GetByCollectionID(inputCollectionID string) ([]*dbm
 	err := s.db.
 		Where("input_collection_id = ?", inputCollectionID).
 		Where("is_deleted = ?", false).
+		Where("is_ready = ?", true).
 		Find(&attachedFunctions).Error
 
 	if err != nil {
@@ -97,7 +100,9 @@ func (s *attachedFunctionDb) GetByCollectionID(inputCollectionID string) ([]*dbm
 func (s *attachedFunctionDb) UpdateOutputCollectionID(id uuid.UUID, outputCollectionID *string) error {
 	now := time.Now()
 	result := s.db.Model(&dbmodel.AttachedFunction{}).
-		Where("id = ? AND is_deleted = false", id).
+		Where("id = ?", id).
+		Where("is_deleted = false").
+		Where("is_ready = true").
 		Updates(map[string]interface{}{
 			"output_collection_id": outputCollectionID,
 			"updated_at":           now,
@@ -164,6 +169,30 @@ func (s *attachedFunctionDb) SoftDeleteByID(id uuid.UUID) error {
 	return nil
 }
 
+// SetReady sets is_ready to true for an attached function
+func (s *attachedFunctionDb) SetReady(id uuid.UUID) error {
+	now := time.Now()
+	result := s.db.Model(&dbmodel.AttachedFunction{}).
+		Where("id = ?", id).
+		Where("is_deleted = false").
+		UpdateColumns(map[string]interface{}{
+			"is_ready":   true,
+			"updated_at": now,
+		})
+
+	if result.Error != nil {
+		log.Error("SetReady failed", zap.Error(result.Error), zap.String("id", id.String()))
+		return result.Error
+	}
+
+	if result.RowsAffected == 0 {
+		log.Error("SetReady: no rows affected - attached function not found", zap.String("id", id.String()))
+		return common.ErrAttachedFunctionNotFound
+	}
+
+	return nil
+}
+
 // Finish marks work as complete
 func (s *attachedFunctionDb) Finish(id uuid.UUID) error {
 	now := time.Now()
@@ -332,3 +361,29 @@ func (s *attachedFunctionDb) HardDeleteAttachedFunction(id uuid.UUID) error {
 
 	return nil
 }
+
+func (s *attachedFunctionDb) UpdateCompletionOffset(id uuid.UUID, completionOffset int64) error {
+	result := s.db.Model(&dbmodel.AttachedFunction{}).
+		Where("id = ? AND is_deleted = ?", id, false).
+		Update("completion_offset", completionOffset)
+
+	if result.Error != nil {
+		log.Error("UpdateCompletionOffset failed",
+			zap.Error(result.Error),
+			zap.String("id", id.String()),
+			zap.Int64("completion_offset", completionOffset))
+		return result.Error
+	}
+
+	if result.RowsAffected == 0 {
+		log.Warn("UpdateCompletionOffset: no rows affected (attached function not found)",
+			zap.String("id", id.String()))
+		return gorm.ErrRecordNotFound
+	}
+
+	log.Debug("UpdateCompletionOffset succeeded",
+		zap.String("id", id.String()),
+		zap.Int64("completion_offset", completionOffset))
+
+	return nil
+}

go/pkg/sysdb/metastore/db/dao/task.go

@@ -26,6 +26,7 @@ type AttachedFunction struct {
 	UpdatedAt               time.Time  `gorm:"column:updated_at;type:timestamp;not null;default:CURRENT_TIMESTAMP"`
 	GlobalParent            *uuid.UUID `gorm:"column:global_parent;type:uuid;default:null"`
 	OldestWrittenNonce      *uuid.UUID `gorm:"column:oldest_written_nonce;type:uuid;default:null"`
+	IsReady                 bool       `gorm:"column:is_ready;type:boolean;not null;default:false"`
 }
 
 func (v AttachedFunction) TableName() string {
@@ -38,6 +39,9 @@ type IAttachedFunctionDb interface {
 	GetByName(inputCollectionID string, name string) (*AttachedFunction, error)
 	GetByID(id uuid.UUID) (*AttachedFunction, error)
 	GetByCollectionID(inputCollectionID string) ([]*AttachedFunction, error)
+	UpdateCompletionOffset(id uuid.UUID, completionOffset int64) error
+	SetReady(id uuid.UUID) error
+	Finish(id uuid.UUID) error
 	UpdateOutputCollectionID(id uuid.UUID, outputCollectionID *string) error
 	SoftDelete(inputCollectionID string, name string) error
 	SoftDeleteByID(id uuid.UUID) error

go/pkg/sysdb/metastore/db/dbmodel/task.go

@@ -0,0 +1,2 @@
+-- Add is_ready column to attached_functions table to track initialization status
+ALTER TABLE "public"."attached_functions" ADD COLUMN "is_ready" boolean NOT NULL DEFAULT false;

go/pkg/sysdb/metastore/db/migrations/20251116154842.sql

@@ -1,4 +1,4 @@
-h1:bIntlbZeBztzLttsjmX5fudfrCbxOd/mgF4fJ91j21s=
+h1:TI+7Y8LdtEPUUsspBjfrE9KdjS8XS8SLfKUoSJhmhEI=
 20240313233558.sql h1:Gv0TiSYsqGoOZ2T2IWvX4BOasauxool8PrBOIjmmIdg=
 20240321194713.sql h1:kVkNpqSFhrXGVGFFvL7JdK3Bw31twFcEhI6A0oCFCkg=
 20240327075032.sql h1:nlr2J74XRU8erzHnKJgMr/tKqJxw9+R6RiiEBuvuzgo=
@@ -25,3 +25,4 @@ h1:bIntlbZeBztzLttsjmX5fudfrCbxOd/mgF4fJ91j21s=
 20251029223300.sql h1:7ptfELfwHoDLLGRz6hLRYAq3OZyzNqa73U5DjlrN/YE=
 20251114125442.sql h1:oRHN+AO+xYnYa3aF0QzSa3T/TRi8ETCydp2sDT/nSnI=
 20251114134400.sql h1:N30qnVNjR+d4RoArJ11YrixyIsNODxqXFpgiREEhczs=
+20251116154842.sql h1:G0qy4MPDayH+Y9/Dm9PS2xwvyLt+nmIw90uJTzZSJUM=