Make sure we wait for the final drain timeout

If the checkDrainTimer fires, we exit immediately even if the
drain has not finished yet. But we should wait for the bigger
timeout as well.

client/cherami/connection.go

@@ -425,18 +425,21 @@ func (conn *connection) waitForDrainAndClose() {
 
 	checkDrainTimer := time.NewTimer(defaultCheckDrainTimeout)
 	defer checkDrainTimer.Stop()
-	select {
-	case <-conn.closeCh:
-		// already closed
-		return
-	case <-checkDrainTimer.C:
-		// check if we got the acks for all sent messages
-		if conn.isAlreadyDrained() {
-			conn.logger.Infof("Inputhost connection drained completely")
+	for {
+		checkDrainTimer.Reset(defaultCheckDrainTimeout)
+		select {
+		case <-conn.closeCh:
+			// already closed
+			return
+		case <-checkDrainTimer.C:
+			// check if we got the acks for all sent messages
+			if conn.isAlreadyDrained() {
+				conn.logger.Infof("Inputhost connection drained completely")
+				return
+			}
+		case <-drainTimer.C:
 			return
 		}
-	case <-drainTimer.C:
-		return
 	}
 }