Improve solver channel management according to @aluzzardi comments

Signed-off-by: Tom Chauveau <tom.chauveau@epitech.eu>

client/client.go

@@ -170,27 +170,18 @@ func (c *Client) buildfn(ctx context.Context, st *state.State, env *environment.
 	go catchOutput(buildCh)
 
 	resp, err := c.c.Build(ctx, opts, "", func(ctx context.Context, gw bkgw.Client) (*bkgw.Result, error) {
-		eventsWg := sync.WaitGroup{}
-		closeCh := make(chan *bk.SolveStatus)
-
-		// Close events channel
-		defer func() {
-			close(closeCh)
-			eventsWg.Wait()
-			close(eventsCh)
-		}()
-
 		s := solver.New(solver.Opts{
 			Control: c.c,
 			Gateway: gw,
 			Events:  eventsCh,
-			EventsWg:   &eventsWg,
-			CloseEvent: closeCh,
 			Auth:    auth,
 			Secrets: secrets,
 			NoCache: c.cfg.NoCache,
 		})
 
+		// Close events channel
+		defer s.Stop()
+
 		// Compute output overlay
 		if fn != nil {
 			if err := fn(ctx, env, s); err != nil {

solver/solver.go

@@ -21,14 +21,14 @@ import (
 
 type Solver struct {
 	opts     Opts
+	eventsWg *sync.WaitGroup
+	closeCh  chan *bk.SolveStatus
 }
 
 type Opts struct {
 	Control *bk.Client
 	Gateway bkgw.Client
 	Events  chan *bk.SolveStatus
-	EventsWg   *sync.WaitGroup
-	CloseEvent chan *bk.SolveStatus
 	Auth    *RegistryAuthProvider
 	Secrets session.Attachable
 	NoCache bool
@@ -36,6 +36,8 @@ type Opts struct {
 
 func New(opts Opts) Solver {
 	return Solver{
+		eventsWg: &sync.WaitGroup{},
+		closeCh:  make(chan *bk.SolveStatus),
 		opts:     opts,
 	}
 }
@@ -63,6 +65,12 @@ func (s Solver) NoCache() bool {
 	return s.opts.NoCache
 }
 
+func (s Solver) Stop() {
+	close(s.closeCh)
+	s.eventsWg.Wait()
+	close(s.opts.Events)
+}
+
 func (s Solver) AddCredentials(target, username, secret string) {
 	s.opts.Auth.AddCredentials(target, username, secret)
 }
@@ -148,11 +156,30 @@ func (s Solver) Solve(ctx context.Context, st llb.State) (bkgw.Reference, error)
 	return res.SingleRef()
 }
 
+// Forward events from solver to the main events channel
+// It creates a task in the solver waiting group to be
+// sure that everything will be forward to the main channel
+func (s Solver) forwardEvents(ch chan *bk.SolveStatus) {
+	s.eventsWg.Add(1)
+	defer s.eventsWg.Done()
+
+	for event := range ch {
+		s.opts.Events <- event
+	}
+}
+
 // Export will export `st` to `output`
 // FIXME: this is currently impleneted as a hack, starting a new Build session
 // within buildkit from the Control API. Ideally the Gateway API should allow to
 // Export directly.
 func (s Solver) Export(ctx context.Context, st llb.State, img *dockerfile2llb.Image, output bk.ExportEntry) (*bk.SolveResponse, error) {
+	// Check close event channel and return if we're already done with the main pipeline
+	select {
+	case <-s.closeCh:
+		return nil, context.Canceled
+	default:
+	}
+
 	def, err := s.Marshal(ctx, st)
 	if err != nil {
 		return nil, err
@@ -171,24 +198,7 @@ func (s Solver) Export(ctx context.Context, st llb.State, img *dockerfile2llb.Im
 
 	// Forward this build session events to the main events channel, for logging
 	// purposes.
-	go func() {
-		select {
-		case <-s.opts.CloseEvent:
-			return
-		default:
-			for event := range ch {
-				s.opts.Events <- event
-			}
-		}
-	}()
-
-	// Add task to events
-	s.opts.EventsWg.Add(1)
-
-	// Resolve event
-	defer func() {
-		s.opts.EventsWg.Done()
-	}()
+	go s.forwardEvents(ch)
 
 	return s.opts.Control.Build(ctx, opts, "", func(ctx context.Context, c bkgw.Client) (*bkgw.Result, error) {
 		res, err := c.Solve(ctx, bkgw.SolveRequest{