1 files changed, 47 insertions, 0 deletions
diff --git a/cpp/src/IceGrid/SessionI.cpp b/cpp/src/IceGrid/SessionI.cpp
index 6afba2bbf10..f7f656e846f 100644
--- a/cpp/src/IceGrid/SessionI.cpp
+++ b/cpp/src/IceGrid/SessionI.cpp
@@ -62,6 +62,43 @@ newAllocateObject(const SessionIPtr& session, const IceUtil::Handle<T>& cb)
     return new AllocateObject<T>(session, cb);
 }
 
+class ConnectionCallbackI : public Ice::ConnectionCallback
+{
+public:
+
+    ConnectionCallbackI(const BaseSessionIPtr& session) : _session(session)
+    {
+    }
+    
+    virtual void
+    heartbeat(const Ice::ConnectionPtr&)
+    {
+        try
+        {
+            _session->keepAlive(Ice::Current());
+        }
+        catch(const Ice::Exception&)
+        {
+        }
+    }
+
+    virtual void
+    closed(const Ice::ConnectionPtr&)
+    {
+        try
+        {
+            _session->destroyImpl(false);
+        }
+        catch(const Ice::Exception&)
+        {
+        }
+    }
+
+private:
+    
+    const BaseSessionIPtr _session;
+};
+
 }
 
 BaseSessionI::BaseSessionI(const string& id, const string& prefix, const DatabasePtr& database) :
@@ -84,6 +121,12 @@ BaseSessionI::~BaseSessionI()
 }
 
 void
+BaseSessionI::setConnectionCallback(const Ice::ConnectionPtr& con)
+{
+    con->setCallback(new ConnectionCallbackI(this));
+}
+
+void
 BaseSessionI::keepAlive(const Ice::Current& current)
 {
     Lock sync(*this);
@@ -169,6 +212,10 @@ SessionI::_register(const SessionServantManagerPtr& servantManager, const Ice::C
     // This is supposed to be called after creation only, no need to synchronize.
     //
     _servantManager = servantManager;
+    if(con)
+    {
+        setConnectionCallback(con);
+    }
     return _servantManager->addSession(this, con, "");
 }