fix: Add shutdown for server

vdaas · Dec 17, 2024 · e9db705 · e9db705
1 parent 1d9062c
commit e9db705
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/pkg/index/job/exportation/service/exporter.go b/pkg/index/job/exportation/service/exporter.go
@@ -46,6 +46,7 @@ const (
 type Exporter interface {
 	StartClient(ctx context.Context) (<-chan error, error)
 	Start(ctx context.Context) error
+	PreStop(ctx context.Context) error
 }
 
 type export struct {
@@ -215,3 +216,7 @@ func (e *export) doExportIndex(
 		}
 	}
 }
+
+func (e *export) PreStop(ctx context.Context) error {
+	return e.storedVector.Close(false)
+}
diff --git a/pkg/index/job/exportation/usecase/exportation.go b/pkg/index/job/exportation/usecase/exportation.go
@@ -116,7 +116,7 @@ func (r *run) PreStart(ctx context.Context) error {
 // during the operation and an error representing any initialization errors.
 func (r *run) Start(ctx context.Context) (<-chan error, error) {
 	ech := make(chan error, 3)
-	var sech, oech, cech <-chan error
+	var sech, oech <-chan error
 	if r.observability != nil {
 		oech = r.observability.Start(ctx)
 	}
@@ -167,8 +167,8 @@ func (r *run) Start(ctx context.Context) (<-chan error, error) {
 }
 
 // PreStop is a method called before execution of Stop.
-func (*run) PreStop(_ context.Context) error {
-	return nil
+func (r *run) PreStop(ctx context.Context) error {
+	return r.exporter.PreStop(ctx)
 }
 
 // Stop is a method used to stop an operation in the run.
@@ -178,6 +178,11 @@ func (r *run) Stop(ctx context.Context) (errs error) {
 			errs = errors.Join(errs, err)
 		}
 	}
+	if r.server != nil {
+		if err := r.server.Shutdown(ctx); err != nil {
+			errs = errors.Join(errs, err)
+		}
+	}
 	return errs
 }