From b11cd08b998093437300df504c8369172899dfcc Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Thu, 18 May 2023 22:01:33 -0700
Subject: [PATCH] session: fix run and close synchronization

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit 64580e7bc084fafb405ff19f0ccb2c097e951698)
Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 session/session.go | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/session/session.go b/session/session.go
index 50cb3b448619..f56a18730d22 100644
--- a/session/session.go
+++ b/session/session.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"net"
 	"strings"
+	"sync"
 
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/moby/buildkit/identity"
@@ -36,14 +37,16 @@ type Attachable interface {
 
 // Session is a long running connection between client and a daemon
 type Session struct {
-	id         string
-	name       string
-	sharedKey  string
-	ctx        context.Context
-	cancelCtx  func()
-	done       chan struct{}
-	grpcServer *grpc.Server
-	conn       net.Conn
+	mu          sync.Mutex // synchronizes conn run and close
+	id          string
+	name        string
+	sharedKey   string
+	ctx         context.Context
+	cancelCtx   func()
+	done        chan struct{}
+	grpcServer  *grpc.Server
+	conn        net.Conn
+	closeCalled bool
 }
 
 // NewSession returns a new long running session
@@ -99,6 +102,11 @@ func (s *Session) ID() string {
 
 // Run activates the session
 func (s *Session) Run(ctx context.Context, dialer Dialer) error {
+	s.mu.Lock()
+	if s.closeCalled {
+		s.mu.Unlock()
+		return nil
+	}
 	ctx, cancel := context.WithCancel(ctx)
 	s.cancelCtx = cancel
 	s.done = make(chan struct{})
@@ -118,15 +126,18 @@ func (s *Session) Run(ctx context.Context, dialer Dialer) error {
 	}
 	conn, err := dialer(ctx, "h2c", meta)
 	if err != nil {
+		s.mu.Unlock()
 		return errors.Wrap(err, "failed to dial gRPC")
 	}
 	s.conn = conn
+	s.mu.Unlock()
 	serve(ctx, s.grpcServer, conn)
 	return nil
 }
 
 // Close closes the session
 func (s *Session) Close() error {
+	s.mu.Lock()
 	if s.cancelCtx != nil && s.done != nil {
 		if s.conn != nil {
 			s.conn.Close()
@@ -134,6 +145,8 @@ func (s *Session) Close() error {
 		s.grpcServer.Stop()
 		<-s.done
 	}
+	s.closeCalled = true
+	s.mu.Unlock()
 	return nil
 }