MB-71131: Support file writer callbacks for trained index (#2379)

- implement the `fileWriterIDs` and `removeFileWriterIDs` APIs which
fetch and remove the file writer IDs/keys corresponding to the trained
segment.
- refactor the request handling + ack mechanism of the trainLoop()
- update the destination index's bolt with the trained sample count, for
consistent stats

Author: Thejas-bhat

index/scorch/scorch.go

@@ -123,6 +123,9 @@ type trainer interface {
 	// trainer specific file transfer operations
 	copyFileLOCKED(file string, d index.IndexDirectory) error
 	updateBolt(snapshotsBucket *util.BoltBucketImpl, key []byte, value []byte) error
+
+	dropFileWriterIDs(ids map[string]struct{}) error
+	fileWriterIDsInUse() (map[string]struct{}, error)
 }
 
 type ScorchErrorType string
@@ -1116,10 +1119,12 @@ func (s *Scorch) SetPathInBolt(key []byte, value []byte) error {
 		return err
 	}
 
-	// currently this is specific to trained index file update
-	err = s.trainer.updateBolt(snapshotsBucket, key, value)
-	if err != nil {
-		return err
+	if s.trainer != nil {
+		// currently this is specific to trained index file update
+		err = s.trainer.updateBolt(snapshotsBucket, key, value)
+		if err != nil {
+			return err
+		}
 	}
 
 	err = tx.Commit()
@@ -1332,6 +1337,16 @@ func (s *Scorch) FileWriterIDsInUse() (map[string]struct{}, error) {
 		}
 	}
 
+	if s.trainer != nil {
+		trainerKeys, err := s.trainer.fileWriterIDsInUse()
+		if err != nil {
+			return nil, err
+		}
+		for k, _ := range trainerKeys {
+			keyMap[k] = struct{}{}
+		}
+	}
+
 	boltKeys, err := s.boltFileWriterIDsInUse()
 	if err != nil {
 		return nil, err
@@ -1354,6 +1369,13 @@ func (s *Scorch) DropFileWriterIDs(ids map[string]struct{}) error {
 		return err
 	}
 
+	if s.trainer != nil {
+		err := s.trainer.dropFileWriterIDs(ids)
+		if err != nil {
+			return err
+		}
+	}
+
 	s.rootLock.Lock()
 	// create a done channel to ensure success of merge
 	ctx := context.Background()

index/scorch/train_noop.go

@@ -53,3 +53,11 @@ func (t *noopTrainer) copyFileLOCKED(file string, d index.IndexDirectory) error
 func (t *noopTrainer) updateBolt(snapshotsBucket *util.BoltBucketImpl, key []byte, value []byte) error {
 	return nil
 }
+
+func (t *noopTrainer) dropFileWriterIDs(ids map[string]struct{}) error {
+	return nil
+}
+
+func (t *noopTrainer) fileWriterIDsInUse() (map[string]struct{}, error) {
+	return nil, nil
+}