Improve disconnect and shutdown handling.

Previously the exported Disconnect and Shutdown functions called each other
and therefore needed to release and reacquire the locks which could
potentionally allow a request to sneak in between the lock and unlock.

This commit changes the exported Shutdown function so that everything is
done under the request lock to prevent this possibility.  In order to
support this, the shared code between the Disconnect and Shutdown
functions has been refactored into unexported functions and the locking
has been altered accordingly.

Also, replace the sendMessage function with a select over the send and
disconnect channels to prevent any issues with queued up messages.

joint debugging effort between myself and @jrick
This commit is contained in:
Dave Collins 2014-07-07 20:11:01 -05:00
parent 1a23feb53e
commit 656fa8699b

View file

@ -145,8 +145,8 @@ func (c *Client) NextID() uint64 {
} }
// addRequest associates the passed jsonRequest with the passed id. This allows // addRequest associates the passed jsonRequest with the passed id. This allows
// the response from the remote server to be unmarshalled to the appropiate type // the response from the remote server to be unmarshalled to the appropriate
// and sent to the specified channel when it is received. // type and sent to the specified channel when it is received.
// //
// If the client has already begun shutting down, ErrClientShutdown is returned // If the client has already begun shutting down, ErrClientShutdown is returned
// and the request is not added. // and the request is not added.
@ -196,11 +196,8 @@ func (c *Client) removeRequest(id uint64) *jsonRequest {
// removeAllRequests removes all the jsonRequests which contain the response // removeAllRequests removes all the jsonRequests which contain the response
// channels for outstanding requests. // channels for outstanding requests.
// //
// This function is safe for concurrent access. // This function MUST be called with the request lock held.
func (c *Client) removeAllRequests() { func (c *Client) removeAllRequests() {
c.requestLock.Lock()
defer c.requestLock.Unlock()
c.requestMap = make(map[uint64]*list.Element) c.requestMap = make(map[uint64]*list.Element)
c.requestList.Init() c.requestList.Init()
} }
@ -416,11 +413,11 @@ cleanup:
// block until the send channel is full. // block until the send channel is full.
func (c *Client) sendMessage(marshalledJSON []byte) { func (c *Client) sendMessage(marshalledJSON []byte) {
// Don't send the message if disconnected. // Don't send the message if disconnected.
if c.Disconnected() { select {
case c.sendChan <- marshalledJSON:
case <-c.disconnect:
return return
} }
c.sendChan <- marshalledJSON
} }
// reregisterNtfns creates and sends commands needed to re-establish the current // reregisterNtfns creates and sends commands needed to re-establish the current
@ -824,14 +821,14 @@ func (c *Client) Disconnected() bool {
return c.disconnected return c.disconnected
} }
// Disconnect disconnects the current websocket associated with the client. The // doDisconnect disconnects the websocket associated with the client if it
// connection will automatically be re-established unless the client was // hasn't already been disconnected. It will return false if the disconnect is
// created with the DisableAutoReconnect flag. // not needed or the client is running in HTTP POST mode.
// //
// This function has no effect when the client is running in HTTP POST mode. // This function is safe for concurrent access.
func (c *Client) Disconnect() { func (c *Client) doDisconnect() bool {
if c.config.HttpPostMode { if c.config.HttpPostMode {
return return false
} }
c.mtx.Lock() c.mtx.Lock()
@ -839,18 +836,51 @@ func (c *Client) Disconnect() {
// Nothing to do if already disconnected. // Nothing to do if already disconnected.
if c.disconnected { if c.disconnected {
return return false
} }
log.Tracef("Disconnecting RPC client %s", c.config.Host) log.Tracef("Disconnecting RPC client %s", c.config.Host)
close(c.disconnect) close(c.disconnect)
c.wsConn.Close() c.wsConn.Close()
c.disconnected = true c.disconnected = true
return true
}
// doShutdown closes the shutdown channel and logs the shutdown unless shutdown
// is already in progress. It will return false if the shutdown is not needed.
//
// This function is safe for concurrent access.
func (c *Client) doShutdown() bool {
// Ignore the shutdown request if the client is already in the process
// of shutting down or already shutdown.
select {
case <-c.shutdown:
return false
default:
}
log.Tracef("Shutting down RPC client %s", c.config.Host)
close(c.shutdown)
return true
}
// Disconnect disconnects the current websocket associated with the client. The
// connection will automatically be re-established unless the client was
// created with the DisableAutoReconnect flag.
//
// This function has no effect when the client is running in HTTP POST mode.
func (c *Client) Disconnect() {
// Nothing to do if already disconnected or running in HTTP POST mode.
if !c.doDisconnect() {
return
}
c.requestLock.Lock()
defer c.requestLock.Unlock()
// When operating without auto reconnect, send errors to any pending // When operating without auto reconnect, send errors to any pending
// requests and shutdown the client. // requests and shutdown the client.
if c.config.DisableAutoReconnect { if c.config.DisableAutoReconnect {
c.requestLock.Lock()
for e := c.requestList.Front(); e != nil; e = e.Next() { for e := c.requestList.Front(); e != nil; e = e.Next() {
req := e.Value.(*jsonRequest) req := e.Value.(*jsonRequest)
req.responseChan <- &response{ req.responseChan <- &response{
@ -858,9 +888,8 @@ func (c *Client) Disconnect() {
err: ErrClientDisconnect, err: ErrClientDisconnect,
} }
} }
c.requestLock.Unlock()
c.removeAllRequests() c.removeAllRequests()
c.Shutdown() c.doShutdown()
} }
} }
@ -868,19 +897,18 @@ func (c *Client) Disconnect() {
// with the client and, when automatic reconnect is enabled, preventing future // with the client and, when automatic reconnect is enabled, preventing future
// attempts to reconnect. It also stops all goroutines. // attempts to reconnect. It also stops all goroutines.
func (c *Client) Shutdown() { func (c *Client) Shutdown() {
// Do the shutdown under the request lock to prevent clients from
// adding new requests while the client shutdown process is initiated.
c.requestLock.Lock()
defer c.requestLock.Unlock()
// Ignore the shutdown request if the client is already in the process // Ignore the shutdown request if the client is already in the process
// of shutting down or already shutdown. // of shutting down or already shutdown.
select { if !c.doShutdown() {
case <-c.shutdown:
return return
default:
} }
log.Tracef("Shutting down RPC client %s", c.config.Host)
close(c.shutdown)
// Send the ErrClientShutdown error to any pending requests. // Send the ErrClientShutdown error to any pending requests.
c.requestLock.Lock()
for e := c.requestList.Front(); e != nil; e = e.Next() { for e := c.requestList.Front(); e != nil; e = e.Next() {
req := e.Value.(*jsonRequest) req := e.Value.(*jsonRequest)
req.responseChan <- &response{ req.responseChan <- &response{
@ -888,10 +916,10 @@ func (c *Client) Shutdown() {
err: ErrClientShutdown, err: ErrClientShutdown,
} }
} }
c.requestLock.Unlock()
c.removeAllRequests() c.removeAllRequests()
c.Disconnect() // Disconnect the client if needed.
c.doDisconnect()
} }
// Start begins processing input and output messages. // Start begins processing input and output messages.