cleanup trainer after complete

Author: Thejas-bhat

index/scorch/snapshot_segment.go

@@ -370,6 +370,9 @@ func (c *cachedMeta) updateMeta(field string, val interface{}) {
 
 func (c *cachedMeta) fetchMeta(field string) (rv interface{}) {
 	c.m.RLock()
+	if c.meta == nil {
+		return nil
+	}
 	rv = c.meta[field]
 	c.m.RUnlock()
 	return rv

index/scorch/train_vector.go

@@ -93,7 +93,10 @@ func (t *vectorTrainer) trainLoop() {
 		t.parent.asyncTasks.Done()
 	}()
 	// initialize stuff
-	totalSamplesProcessed := t.centroidIndex.cachedMeta.fetchMeta("trainedSamples").(uint64)
+	var totalSamplesProcessed uint64
+	if t.centroidIndex != nil {
+		totalSamplesProcessed = t.centroidIndex.cachedMeta.fetchMeta("trainedSamples").(uint64)
+	}
 	buf := make([]byte, binary.MaxVarintLen64)
 	t.parent.segmentConfig[index.CentroidIndexCallback] = t.getCentroidIndex
 	path := filepath.Join(t.parent.path, index.CentroidIndexFileName)
@@ -230,9 +233,23 @@ func (t *vectorTrainer) trainLoop() {
 			}
 			t.m.Lock()
 			t.centroidIndex = &SegmentSnapshot{
-				segment: centroidIndex,
+				segment:    centroidIndex,
+				cachedMeta: &cachedMeta{meta: nil},
 			}
 			t.m.Unlock()
+
+			// if the train complete flag has been set, exit the routine and cleanup
+			if trainReq.trainComplete {
+				// cleanup .tmp file if it exists
+				if _, err := os.Stat(path + ".tmp"); err == nil {
+					err = os.Remove(path + ".tmp")
+					if err != nil {
+						trainReq.ackCh <- fmt.Errorf("error removing .tmp file: %v", err)
+					}
+				}
+				close(trainReq.ackCh)
+				return
+			}
 			close(trainReq.ackCh)
 		}
 	}

index_impl.go

@@ -1465,38 +1465,6 @@ func (i *indexImpl) CopyTo(d index.Directory) (err error) {
 	return i.meta.CopyTo(d)
 }
 
-func (i *indexImpl) CopyFile(file string, d index.IndexDirectory) (err error) {
-	i.mutex.RLock()
-	defer i.mutex.RUnlock()
-
-	if !i.open {
-		return ErrorIndexClosed
-	}
-
-	copyIndex, ok := i.i.(index.IndexFileCopyable)
-	if !ok {
-		return fmt.Errorf("index implementation does not support copy reader")
-	}
-
-	return copyIndex.CopyFile(file, d)
-}
-
-func (i *indexImpl) UpdateFileInBolt(key []byte, value []byte) error {
-	i.mutex.RLock()
-	defer i.mutex.RUnlock()
-
-	if !i.open {
-		return ErrorIndexClosed
-	}
-
-	copyIndex, ok := i.i.(index.IndexFileCopyable)
-	if !ok {
-		return fmt.Errorf("index implementation does not support file copy")
-	}
-
-	return copyIndex.UpdateFileInBolt(key, value)
-}
-
 func (f FileSystemDirectory) GetWriter(filePath string) (io.WriteCloser,
 	error,
 ) {