2 files changed, 9 insertions, 4 deletions
diff --git a/Swiften/Network/BoostConnection.cpp b/Swiften/Network/BoostConnection.cpp
index 043743e..534ebdb 100644
--- a/Swiften/Network/BoostConnection.cpp
+++ b/Swiften/Network/BoostConnection.cpp
@@ -19,105 +19,109 @@
 #include <Swiften/Base/Algorithm.h>
 #include <Swiften/EventLoop/EventLoop.h>
 #include <Swiften/Base/ByteArray.h>
 #include <Swiften/Network/HostAddressPort.h>
 #include <Swiften/Base/sleep.h>
 #include <Swiften/Base/SafeAllocator.h>
 
 namespace Swift {
 
 static const size_t BUFFER_SIZE = 4096;
 
 // -----------------------------------------------------------------------------
 
 // A reference-counted non-modifiable buffer class.
 class SharedBuffer {
 	public:
 		SharedBuffer(const SafeByteArray& data) :
 				data_(new std::vector<char, SafeAllocator<char> >(data.begin(), data.end())),
 				buffer_(boost::asio::buffer(*data_)) {
 		}
 
 		// ConstBufferSequence requirements.
 		typedef boost::asio::const_buffer value_type;
 		typedef const boost::asio::const_buffer* const_iterator;
 		const boost::asio::const_buffer* begin() const { return &buffer_; }
 		const boost::asio::const_buffer* end() const { return &buffer_ + 1; }
 
 	private:
 		boost::shared_ptr< std::vector<char, SafeAllocator<char> > > data_;
 		boost::asio::const_buffer buffer_;
 };
 
 // -----------------------------------------------------------------------------
 
 BoostConnection::BoostConnection(boost::shared_ptr<boost::asio::io_service> ioService, EventLoop* eventLoop) :
-		eventLoop(eventLoop), ioService(ioService), socket_(*ioService), writing_(false) {
+	eventLoop(eventLoop), ioService(ioService), socket_(*ioService), writing_(false), closeSocketAfterNextWrite_(false) {
 }
 
 BoostConnection::~BoostConnection() {
 }
 
 void BoostConnection::listen() {
 	doRead();
 }
 
 void BoostConnection::connect(const HostAddressPort& addressPort) {
 	boost::asio::ip::tcp::endpoint endpoint(	
 			boost::asio::ip::address::from_string(addressPort.getAddress().toString()), addressPort.getPort());
 	socket_.async_connect(
 			endpoint,
 			boost::bind(&BoostConnection::handleConnectFinished, shared_from_this(), boost::asio::placeholders::error));
 }
 
 void BoostConnection::disconnect() {
 	//MainEventLoop::removeEventsFromOwner(shared_from_this());
 
 	// Mac OS X apparently exhibits a problem where closing a socket during a write could potentially go into uninterruptable sleep.
 	// See e.g. http://bugs.python.org/issue7401
 	// We therefore wait until any pending write finishes, which hopefully should fix our hang on exit during close().
-	while (writing_) {
-		Swift::sleep(10);
+	if (writing_) {
+		closeSocketAfterNextWrite_ = true;
+	} else {
+		socket_.close();
 	}
-	socket_.close();
 }
 
 void BoostConnection::write(const SafeByteArray& data) {
 	boost::lock_guard<boost::mutex> lock(writeMutex_);
 	if (!writing_) {
 		writing_ = true;
 		doWrite(data);
+		if (closeSocketAfterNextWrite_) {
+			socket_.close();
+		}
 	}
 	else {
 		append(writeQueue_, data);
 	}
 }
 
 void BoostConnection::doWrite(const SafeByteArray& data) {
 	boost::asio::async_write(socket_, SharedBuffer(data),
 			boost::bind(&BoostConnection::handleDataWritten, shared_from_this(), boost::asio::placeholders::error));
 }
 
 void BoostConnection::handleConnectFinished(const boost::system::error_code& error) {
 	SWIFT_LOG(debug) << "Connect finished: " << error << std::endl;
 	if (!error) {
 		eventLoop->postEvent(boost::bind(boost::ref(onConnectFinished), false), shared_from_this());
 		doRead();
 	}
 	else if (error != boost::asio::error::operation_aborted) {
 		eventLoop->postEvent(boost::bind(boost::ref(onConnectFinished), true), shared_from_this());
 	}
 }
 
 void BoostConnection::doRead() {
 	readBuffer_ = boost::make_shared<SafeByteArray>(BUFFER_SIZE);
 	socket_.async_read_some(
 			boost::asio::buffer(*readBuffer_),
 			boost::bind(&BoostConnection::handleSocketRead, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
 }
 
 void BoostConnection::handleSocketRead(const boost::system::error_code& error, size_t bytesTransferred) {
 	SWIFT_LOG(debug) << "Socket read " << error << std::endl;
 	if (!error) {
 		readBuffer_->resize(bytesTransferred);
 		eventLoop->postEvent(boost::bind(boost::ref(onDataRead), readBuffer_), shared_from_this());
 		doRead();
diff --git a/Swiften/Network/BoostConnection.h b/Swiften/Network/BoostConnection.h
index 7d5ec60..2f0c7be 100644
--- a/Swiften/Network/BoostConnection.h
+++ b/Swiften/Network/BoostConnection.h
@@ -31,37 +31,38 @@ namespace Swift {
 
 			~BoostConnection();
 
 			static ref create(boost::shared_ptr<boost::asio::io_service> ioService, EventLoop* eventLoop) {
 				return ref(new BoostConnection(ioService, eventLoop));
 			}
 
 			virtual void listen();
 			virtual void connect(const HostAddressPort& address);
 			virtual void disconnect();
 			virtual void write(const SafeByteArray& data);
 
 			boost::asio::ip::tcp::socket& getSocket() {
 				return socket_;
 			}
 
 			HostAddressPort getLocalAddress() const;
 
 		private:
 			BoostConnection(boost::shared_ptr<boost::asio::io_service> ioService, EventLoop* eventLoop);
 
 			void handleConnectFinished(const boost::system::error_code& error);
 			void handleSocketRead(const boost::system::error_code& error, size_t bytesTransferred);
 			void handleDataWritten(const boost::system::error_code& error);
 			void doRead();
 			void doWrite(const SafeByteArray& data);
 
 		private:
 			EventLoop* eventLoop;
 			boost::shared_ptr<boost::asio::io_service> ioService;
 			boost::asio::ip::tcp::socket socket_;
 			boost::shared_ptr<SafeByteArray> readBuffer_;
 			boost::mutex writeMutex_;
 			bool writing_;
 			SafeByteArray writeQueue_;
+			bool closeSocketAfterNextWrite_;
 	};
 }