fixed some bugs in socket and stream RpcChannels

ersintarhan · Jun 17, 2010 · e83c62f · e83c62f
1 parent 8296cac
commit e83c62f
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 8 deletions.
diff --git a/Cpp/ProtoBufRemote/RpcServer.cpp b/Cpp/ProtoBufRemote/RpcServer.cpp
@@ -27,6 +27,14 @@ bool RpcServer::UnregisterService(RpcService* service)
 	return (m_services.erase(service->GetName()) == 1);
 }
 
+RpcService* RpcServer::FindService(const std::string& serviceName)
+{
+	ServiceMap::iterator iter = m_services.find(serviceName);
+	if (iter != m_services.end())
+		return iter->second;
+	return NULL;
+}
+
 void RpcServer::ReceiveCall(const RpcMessage& message)
 {
 	const RpcMessage::Call& callMessage = message.call_message();

diff --git a/Cpp/ProtoBufRemote/RpcServer.h b/Cpp/ProtoBufRemote/RpcServer.h
@@ -18,6 +18,8 @@ class RpcServer
 	bool RegisterService(RpcService* service);
 	bool UnregisterService(RpcService* service);
 
+	RpcService* FindService(const std::string& serviceName);
+
 	virtual void ReceiveCall(const RpcMessage& message);
 
 private:

diff --git a/Cpp/ProtoBufRemote/SocketRpcChannel.cpp b/Cpp/ProtoBufRemote/SocketRpcChannel.cpp
@@ -63,7 +63,8 @@ void SocketRpcChannel::Run()
 	const unsigned int maxAllowedMessageSize = 1024*1024;
 
 	WSAEVENT events[3] = { m_terminateEvent, selectEvent, m_sendEvent };
-	for ( ;; )
+	bool isTerminated = false;
+	while (!isTerminated)
 	{
 		int waitResult = WSAWaitForMultipleEvents(isSending ? 2 : 3, events, FALSE, WSA_INFINITE, FALSE);
 		if (waitResult == 0)
@@ -92,10 +93,18 @@ void SocketRpcChannel::Run()
 			isReceiveReady = (networkEvents.lNetworkEvents & FD_READ) ? true : false;
 		}
 
-		if (isSending && isSendReady)
+		while (isSending && isSendReady && !isTerminated)
 		{
 			int bytesSent = send(m_socket, &currentSendMessage.m_data[sendPos], currentSendMessage.m_size-sendPos, 0);
-			if (bytesSent != SOCKET_ERROR)
+			if (bytesSent == SOCKET_ERROR)
+			{
+				int error = WSAGetLastError();
+				if (error == WSAEWOULDBLOCK)
+					isSendReady = false;
+				else
+					isTerminated = true;
+			}
+			else
 			{
 				sendPos += bytesSent;
 
@@ -119,12 +128,22 @@ void SocketRpcChannel::Run()
 			}
 		}
 
-		if (isReceiveReady)
+		while (isReceiveReady && !isTerminated)
 		{
 			int bytesReceived = recv(m_socket, &receiveBuffer[receivePos], receiveSize-receivePos, 0);
 			if (bytesReceived == 0)
-				break;
-			if (bytesReceived != SOCKET_ERROR)
+			{
+				isTerminated = true;
+			}
+			else if (bytesReceived == SOCKET_ERROR)
+			{
+				int error = WSAGetLastError();
+				if (error == WSAEWOULDBLOCK)
+					isReceiveReady = false;
+				else
+					isTerminated = true;
+			}
+			else
 			{
 				receivePos += bytesReceived;
 

diff --git a/DotNet/ProtoBufRemote/StreamRpcChannel.cs b/DotNet/ProtoBufRemote/StreamRpcChannel.cs
@@ -79,6 +79,10 @@ private void ReadRun()
                     if (bytesRead == 0)
                         break;
                 }
+                catch (IOException)
+                {
+                    break;
+                }
                 catch (ObjectDisposedException)
                 {
                     break;
@@ -115,15 +119,24 @@ private void ReadRun()
         private void WriteRun()
         {
             var waitHandles = new WaitHandle[] { queueEvent, closeEvent };
-            while (true)
+            bool isTerminated = false;
+            while (!isTerminated)
             {
                 int waitIndex = WaitHandle.WaitAny(waitHandles);
 
                 queueMutex.WaitOne();
                 while (queuedMessages.Count > 0)
                 {
                     RpcMessage message = queuedMessages.Dequeue();
-                    Serializer.SerializeWithLengthPrefix(writeStream, message, PrefixStyle.Fixed32);
+                    try
+                    {
+                        Serializer.SerializeWithLengthPrefix(writeStream, message, PrefixStyle.Fixed32);
+                    }
+                    catch (InvalidOperationException)
+                    {
+                        isTerminated = true;
+                        break;
+                    }
                 }
                 queueMutex.ReleaseMutex();