@@ -731,50 +731,63 @@ func (c *Client) Close() error {
731
731
732
732
// serveSocketMount persistently listens to the socketMounts listener and proxies connections to a
733
733
// 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 {
735
735
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
761
754
}
762
- _ = cConn .Close ()
763
- return
755
+ return err
764
756
}
765
757
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
+ }
769
777
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
+ }
778
791
}
779
792
}
780
793
0 commit comments