@@ -731,50 +731,63 @@ func (c *Client) Close() error {
731731
732732// serveSocketMount persistently listens to the socketMounts listener and proxies connections to a
733733// given Cloud SQL instance.
734- func (c * Client ) serveSocketMount (_ context.Context , s * socketMount ) error {
734+ func (c * Client ) serveSocketMount (ctx context.Context , s * socketMount ) error {
735735 for {
736- cConn , err := s .Accept ()
737- if err != nil {
738- if nerr , ok := err .(net.Error ); ok && nerr .Timeout () {
739- c .logger .Errorf ("[%s] Error accepting connection: %v" , s .inst , err )
740- // For transient errors, wait a small amount of time to see if it resolves itself
741- time .Sleep (10 * time .Millisecond )
742- continue
743- }
744- return err
745- }
746- // handle the connection in a separate goroutine
747- go func () {
748- c .logger .Infof ("[%s] Accepted connection from %s" , s .inst , cConn .RemoteAddr ())
749-
750- // A client has established a connection to the local socket. Before
751- // we initiate a connection to the Cloud SQL backend, increment the
752- // connection counter. If the total number of connections exceeds
753- // the maximum, refuse to connect and close the client connection.
754- count := atomic .AddUint64 (& c .connCount , 1 )
755- defer atomic .AddUint64 (& c .connCount , ^ uint64 (0 ))
756-
757- if c .conf .MaxConnections > 0 && count > c .conf .MaxConnections {
758- c .logger .Infof ("max connections (%v) exceeded, refusing new connection" , c .conf .MaxConnections )
759- if c .connRefuseNotify != nil {
760- go c .connRefuseNotify ()
736+ select {
737+ case <- ctx .Done ():
738+ // If the context was canceled, do not accept any more connections,
739+ // exit gracefully.
740+ return nil
741+ default :
742+ // Wait to accept a connection. When s.Accept() returns io.EOF, exit
743+ // gracefully.
744+ cConn , err := s .Accept ()
745+ if err != nil {
746+ if nerr , ok := err .(net.Error ); ok && nerr .Timeout () {
747+ c .logger .Errorf ("[%s] Error accepting connection: %v" , s .inst , err )
748+ // For transient errors, wait a small amount of time to see if it resolves itself
749+ time .Sleep (10 * time .Millisecond )
750+ continue
751+ } else if err == io .EOF {
752+ // The socket was closed gracefully. Stop processing connections.
753+ return nil
761754 }
762- _ = cConn .Close ()
763- return
755+ return err
764756 }
765757
766- // give a max of 30 seconds to connect to the instance
767- ctx , cancel := context .WithTimeout (context .Background (), 30 * time .Second )
768- defer cancel ()
758+ // handle the connection in a separate goroutine
759+ go func () {
760+ c .logger .Infof ("[%s] Accepted connection from %s" , s .inst , cConn .RemoteAddr ())
761+
762+ // A client has established a connection to the local socket. Before
763+ // we initiate a connection to the Cloud SQL backend, increment the
764+ // connection counter. If the total number of connections exceeds
765+ // the maximum, refuse to connect and close the client connection.
766+ count := atomic .AddUint64 (& c .connCount , 1 )
767+ defer atomic .AddUint64 (& c .connCount , ^ uint64 (0 ))
768+
769+ if c .conf .MaxConnections > 0 && count > c .conf .MaxConnections {
770+ c .logger .Infof ("max connections (%v) exceeded, refusing new connection" , c .conf .MaxConnections )
771+ if c .connRefuseNotify != nil {
772+ go c .connRefuseNotify ()
773+ }
774+ _ = cConn .Close ()
775+ return
776+ }
769777
770- sConn , err := c .dialer .Dial (ctx , s .inst , s .dialOpts ... )
771- if err != nil {
772- c .logger .Errorf ("[%s] failed to connect to instance: %v" , s .inst , err )
773- _ = cConn .Close ()
774- return
775- }
776- c .proxyConn (s .inst , cConn , sConn )
777- }()
778+ // give a max of 30 seconds to connect to the instance
779+ ctx , cancel := context .WithTimeout (context .Background (), 30 * time .Second )
780+ defer cancel ()
781+
782+ sConn , err := c .dialer .Dial (ctx , s .inst , s .dialOpts ... )
783+ if err != nil {
784+ c .logger .Errorf ("[%s] failed to connect to instance: %v" , s .inst , err )
785+ _ = cConn .Close ()
786+ return
787+ }
788+ c .proxyConn (s .inst , cConn , sConn )
789+ }()
790+ }
778791 }
779792}
780793
0 commit comments