Ref tomb

cmd/crowdsec-cli/clinotifications/notifications.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/csv"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"io/fs"
 	"net/url"
@@ -19,7 +18,6 @@ import (
 	"github.com/go-openapi/strfmt"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
-	"gopkg.in/tomb.v2"
 	"gopkg.in/yaml.v3"
 
 	"github.com/crowdsecurity/go-cs-lib/ptr"
@@ -275,7 +273,6 @@ func (cli *cliNotifications) notificationConfigFilter(_ *cobra.Command, args []s
 func (cli *cliNotifications) newTestCmd() *cobra.Command {
 	var (
 		pluginBroker  csplugin.PluginBroker
-		pluginTomb    tomb.Tomb
 		alertOverride string
 	)
 
@@ -315,11 +312,19 @@ func (cli *cliNotifications) newTestCmd() *cobra.Command {
 				},
 			}, cfg.ConfigPaths)
 		},
-		RunE: func(_ *cobra.Command, _ []string) error {
-			pluginTomb.Go(func() error {
-				pluginBroker.Run(&pluginTomb)
-				return nil
-			})
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			pluginCtx, cancelPlugin := context.WithCancel(cmd.Context())
+			pluginDone := make(chan struct{})
+
+			go func() {
+				pluginBroker.Run(pluginCtx)
+				close(pluginDone)
+			}()
+			defer func() {
+				cancelPlugin()
+				<-pluginDone
+			}()
+
 			alert := &models.Alert{
 				Capacity: ptr.Of(int32(0)),
 				Decisions: []*models.Decision{{
@@ -360,16 +365,12 @@ func (cli *cliNotifications) newTestCmd() *cobra.Command {
 				Alert:     alert,
 			}
 
-			// time.Sleep(2 * time.Second) // There's no mechanism to ensure notification has been sent
-			pluginTomb.Kill(errors.New("terminating"))
-			_ = pluginTomb.Wait()
-
 			return nil
 		},
 	}
 	cmd.Flags().StringVarP(&alertOverride, "alert", "a", "",
-		"JSON string used to override alert fields in the generic alert " +
-		"(see crowdsec/pkg/models/alert.go in the source tree for the full definition of the object)")
+		"JSON string used to override alert fields in the generic alert "+
+			"(see crowdsec/pkg/models/alert.go in the source tree for the full definition of the object)")
 
 	return cmd
 }
@@ -401,10 +402,7 @@ cscli notifications reinject <alert_id> -a '{"remediation": true,"scenario":"not
 			return nil
 		},
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			var (
-				pluginBroker csplugin.PluginBroker
-				pluginTomb   tomb.Tomb
-			)
+			var pluginBroker csplugin.PluginBroker
 
 			ctx := cmd.Context()
 			cfg := cli.cfg()
@@ -442,10 +440,17 @@ cscli notifications reinject <alert_id> -a '{"remediation": true,"scenario":"not
 				return fmt.Errorf("can't initialize plugins: %w", err)
 			}
 
-			pluginTomb.Go(func() error {
-				pluginBroker.Run(&pluginTomb)
-				return nil
-			})
+			pluginCtx, cancelPlugin := context.WithCancel(ctx)
+			pluginDone := make(chan struct{})
+
+			go func() {
+				pluginBroker.Run(pluginCtx)
+				close(pluginDone)
+			}()
+			defer func() {
+				cancelPlugin()
+				<-pluginDone
+			}()
 
 			profiles, err := csprofiles.NewProfile(cfg.API.Server.Profiles)
 			if err != nil {
@@ -481,16 +486,12 @@ cscli notifications reinject <alert_id> -a '{"remediation": true,"scenario":"not
 					break
 				}
 			}
-			// time.Sleep(2 * time.Second) // There's no mechanism to ensure notification has been sent
-			pluginTomb.Kill(errors.New("terminating"))
-			_ = pluginTomb.Wait()
-
 			return nil
 		},
 	}
 	cmd.Flags().StringVarP(&alertOverride, "alert", "a", "",
-		"JSON string used to override alert fields in the reinjected alert " +
-		"(see crowdsec/pkg/models/alert.go in the source tree for the full definition of the object)")
+		"JSON string used to override alert fields in the reinjected alert "+
+			"(see crowdsec/pkg/models/alert.go in the source tree for the full definition of the object)")
 
 	return cmd
 }

cmd/crowdsec/api.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"net"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -20,6 +21,15 @@ func initAPIServer(ctx context.Context, cConfig *csconfig.Config) (*apiserver.AP
 		log.Info("push and pull to Central API disabled")
 	}
 
+	if cConfig.API.Server.ListenURI != "" {
+		listenConfig := &net.ListenConfig{}
+		listener, err := listenConfig.Listen(ctx, "tcp", cConfig.API.Server.ListenURI)
+		if err != nil {
+			return nil, fmt.Errorf("local API server stopped with error: listening on %s: %w", cConfig.API.Server.ListenURI, err)
+		}
+		_ = listener.Close()
+	}
+
 	accessLogger := cConfig.API.Server.NewAccessLogger(cConfig.Common.LogConfig, accessLogFilename)
 
 	apiServer, err := apiserver.NewServer(ctx, cConfig.API.Server, accessLogger)
@@ -42,30 +52,62 @@ func initAPIServer(ctx context.Context, cConfig *csconfig.Config) (*apiserver.AP
 
 func serveAPIServer(ctx context.Context, apiServer *apiserver.APIServer) {
 	apiReady := make(chan bool, 1)
+	runCtx, cancel := context.WithCancel(ctx)
+	apiCancel = cancel
+	apiDone = make(chan error, 1)
 
-	apiTomb.Go(func() error {
+	go func() {
 		defer trace.ReportPanic()
 
+		pluginCtx, cancelPlugin := context.WithCancel(runCtx)
+		pluginDone := make(chan struct{})
+
+		go func() {
+			defer trace.ReportPanic()
+			defer close(pluginDone)
+			pluginBroker.Run(pluginCtx)
+		}()
+
+		runErr := make(chan error, 1)
+
 		go func() {
 			defer trace.ReportPanic()
 
 			log.Debugf("serving API after %s ms", time.Since(crowdsecT0))
 
-			if err := apiServer.Run(ctx, apiReady); err != nil {
-				log.Fatal(err)
-			}
+			runErr <- apiServer.Run(runCtx, apiReady)
 		}()
 
-		pluginTomb.Go(func() error {
-			pluginBroker.Run(&pluginTomb)
-			return nil
-		})
-
-		<-apiTomb.Dying() // lock until go routine is dying
-		pluginTomb.Kill(nil)
-		log.Infof("serve: shutting down api server")
+		select {
+		case err := <-runErr:
+			if err != nil && runCtx.Err() == nil {
+				log.Fatal(err)
+			}
+			cancelPlugin()
+			<-pluginDone
+			apiDone <- err
+		case <-runCtx.Done():
+			log.Infof("serve: shutting down api server")
+
+			shutdownCtx, cancelShutdown := context.WithTimeout(context.WithoutCancel(runCtx), 5*time.Second)
+			err := apiServer.Shutdown(shutdownCtx)
+			cancelShutdown()
+
+			if runErrValue := <-runErr; runErrValue != nil && err == nil {
+				err = runErrValue
+			}
 
-		return apiServer.Shutdown(ctx)
-	})
-	<-apiReady
+			cancelPlugin()
+			<-pluginDone
+			apiDone <- err
+		}
+	}()
+
+	select {
+	case <-apiReady:
+	case err := <-apiDone:
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
 }

cmd/crowdsec/crowdsec.go

@@ -185,38 +185,37 @@ func serveCrowdsec(
 	sd *StateDumper,
 ) {
 	cctx, cancel := context.WithCancel(ctx)
+	crowdCancel = cancel
+	crowdDone = make(chan error, 1)
 
 	var g errgroup.Group
 
 	bucketStore := leakybucket.NewBucketStore()
 
-	crowdsecTomb.Go(func() error {
+	go func() {
 		defer trace.ReportPanic()
+		// this logs every time, even at config reload
+		log.Debugf("running agent after %s ms", time.Since(crowdsecT0))
+		agentReady <- true
 
-		go func() {
-			defer trace.ReportPanic()
-			// this logs every time, even at config reload
-			log.Debugf("running agent after %s ms", time.Since(crowdsecT0))
-
-			agentReady <- true
-
-			if err := runCrowdsec(cctx, &g, cConfig, parsers, hub, datasources, sd, bucketStore); err != nil {
-				log.Fatalf("unable to start crowdsec routines: %s", err)
-			}
-		}()
+		if err := runCrowdsec(cctx, &g, cConfig, parsers, hub, datasources, sd, bucketStore); err != nil {
+			crowdDone <- fmt.Errorf("unable to start crowdsec routines: %w", err)
+			return
+		}
 
 		/* we should stop in two cases :
-		- crowdsecTomb has been Killed() : it might be shutdown or reload, so stop
+		- context has been canceled: it might be shutdown or reload, so stop
 		- acquisTomb is dead, it means that we were in "cat" mode and files are done reading, quit
 		*/
-		waitOnTomb()
+		waitOnCrowdsecStop(cctx)
 		log.Debugf("Shutting down crowdsec routines")
 
 		if err := ShutdownCrowdsecRoutines(cancel, &g, datasources); err != nil {
-			return fmt.Errorf("unable to shutdown crowdsec routines: %w", err)
+			crowdDone <- fmt.Errorf("unable to shutdown crowdsec routines: %w", err)
+			return
 		}
 
-		log.Debugf("everything is dead, return crowdsecTomb")
+		log.Debugf("everything is dead, return crowdsec routine")
 		log.Debugf("sd.DumpDir == %s", sd.DumpDir)
 
 		if sd.DumpDir != "" {
@@ -229,11 +228,11 @@ func serveCrowdsec(
 			os.Exit(0)
 		}
 
-		return nil
-	})
+		crowdDone <- nil
+	}()
 }
 
-func waitOnTomb() {
+func waitOnCrowdsecStop(ctx context.Context) {
 	for {
 		select {
 		case <-acquisTomb.Dead():
@@ -256,7 +255,7 @@ func waitOnTomb() {
 
 			return
 
-		case <-crowdsecTomb.Dying():
+		case <-ctx.Done():
 			log.Infof("Crowdsec engine shutting down")
 			return
 		}

cmd/crowdsec/main.go

@@ -30,12 +30,13 @@ import (
 )
 
 var (
-	// tombs for the parser, buckets and outputs.
-	acquisTomb   tomb.Tomb
-	outputsTomb  tomb.Tomb
-	apiTomb      tomb.Tomb
-	crowdsecTomb tomb.Tomb
-	pluginTomb   tomb.Tomb
+	// tombs for acquisition and output routines.
+	acquisTomb  tomb.Tomb
+	outputsTomb tomb.Tomb
+	apiCancel   context.CancelFunc
+	apiDone     chan error
+	crowdCancel context.CancelFunc
+	crowdDone   chan error
 
 	flags Flags