5 files changed, 39 insertions, 3 deletions

diff --git a/src/vnet/session/application.h b/src/vnet/session/application.h
index a17270cf142..17f2f4a076f 100644
--- a/src/vnet/session/application.h
+++ b/src/vnet/session/application.h
@@ -251,6 +251,8 @@ int app_worker_connect_notify (app_worker_t * app_wrk, session_t * s,
 			       u32 opaque);
 int app_worker_close_notify (app_worker_t * app_wrk, session_t * s);
 int app_worker_reset_notify (app_worker_t * app_wrk, session_t * s);
+int app_worker_cleanup_notify (app_worker_t * app_wrk, session_t * s,
+			       session_cleanup_ntf_t ntf);
 int app_worker_builtin_rx (app_worker_t * app_wrk, session_t * s);
 int app_worker_builtin_tx (app_worker_t * app_wrk, session_t * s);
 segment_manager_t *app_worker_get_listen_segment_manager (app_worker_t *,
diff --git a/src/vnet/session/application_interface.h b/src/vnet/session/application_interface.h
index 234813e5814..f5a0ec08b9b 100644
--- a/src/vnet/session/application_interface.h
+++ b/src/vnet/session/application_interface.h
@@ -39,6 +39,9 @@ typedef struct _stream_session_cb_vft
   /** Notify app that session is closing */
   void (*session_disconnect_callback) (session_t * s);
 
+  /** Notify app that session or transport are about to be removed */
+  void (*session_cleanup_callback) (session_t * s, session_cleanup_ntf_t ntf);
+
   /** Notify app that session was reset */
   void (*session_reset_callback) (session_t * s);
 
diff --git a/src/vnet/session/application_worker.c b/src/vnet/session/application_worker.c
index 69b19909b56..84682cdf930 100644
--- a/src/vnet/session/application_worker.c
+++ b/src/vnet/session/application_worker.c
@@ -334,6 +334,16 @@ app_worker_reset_notify (app_worker_t * app_wrk, session_t * s)
 }
 
 int
+app_worker_cleanup_notify (app_worker_t * app_wrk, session_t * s,
+			   session_cleanup_ntf_t ntf)
+{
+  application_t *app = application_get (app_wrk->app_index);
+  if (app->cb_fns.session_cleanup_callback)
+    app->cb_fns.session_cleanup_callback (s, ntf);
+  return 0;
+}
+
+int
 app_worker_builtin_rx (app_worker_t * app_wrk, session_t * s)
 {
   application_t *app = application_get (app_wrk->app_index);
diff --git a/src/vnet/session/session.c b/src/vnet/session/session.c
index d378d6d2426..5b4e840f053 100644
--- a/src/vnet/session/session.c
+++ b/src/vnet/session/session.c
@@ -179,9 +179,21 @@ session_free (session_t * s)
   pool_put (session_main.wrk[s->thread_index].sessions, s);
 }
 
+static void
+session_cleanup_notify (session_t * s, session_cleanup_ntf_t ntf)
+{
+  app_worker_t *app_wrk;
+
+  app_wrk = app_worker_get_if_valid (s->app_wrk_index);
+  if (!app_wrk)
+    return;
+  app_worker_cleanup_notify (app_wrk, s, ntf);
+}
+
 void
 session_free_w_fifos (session_t * s)
 {
+  session_cleanup_notify (s, SESSION_CLEANUP_SESSION);
   segment_manager_dealloc_fifos (s->rx_fifo, s->tx_fifo);
   session_free (s);
 }
@@ -774,9 +786,6 @@ session_transport_delete_notify (transport_connection_t * tc)
   if (!(s = session_get_if_valid (tc->s_index, tc->thread_index)))
     return;
 
-  /* Make sure we don't try to send anything more */
-  svm_fifo_dequeue_drop_all (s->tx_fifo);
-
   switch (s->session_state)
     {
     case SESSION_STATE_CREATED:
@@ -794,6 +803,8 @@ session_transport_delete_notify (transport_connection_t * tc)
        * are assumed to have been removed from the lookup table */
       session_lookup_del_session (s);
       s->session_state = SESSION_STATE_TRANSPORT_CLOSED;
+      session_cleanup_notify (s, SESSION_CLEANUP_TRANSPORT);
+      svm_fifo_dequeue_drop_all (s->tx_fifo);
       break;
     case SESSION_STATE_CLOSING:
     case SESSION_STATE_CLOSED_WAITING:
@@ -805,14 +816,18 @@ session_transport_delete_notify (transport_connection_t * tc)
       session_lookup_del_session (s);
       s->session_state = SESSION_STATE_TRANSPORT_CLOSED;
       session_program_transport_close (s);
+      session_cleanup_notify (s, SESSION_CLEANUP_TRANSPORT);
+      svm_fifo_dequeue_drop_all (s->tx_fifo);
       break;
     case SESSION_STATE_TRANSPORT_CLOSED:
       break;
     case SESSION_STATE_CLOSED:
+      session_cleanup_notify (s, SESSION_CLEANUP_TRANSPORT);
       session_delete (s);
       break;
     default:
       clib_warning ("session state %u", s->session_state);
+      session_cleanup_notify (s, SESSION_CLEANUP_TRANSPORT);
       session_delete (s);
       break;
     }
diff --git a/src/vnet/session/session_types.h b/src/vnet/session/session_types.h
index be1111f3b20..3564ee77377 100644
--- a/src/vnet/session/session_types.h
+++ b/src/vnet/session/session_types.h
@@ -108,6 +108,12 @@ session_endpoint_is_zero (session_endpoint_t * sep)
 typedef u8 session_type_t;
 typedef u64 session_handle_t;
 
+typedef enum
+{
+  SESSION_CLEANUP_TRANSPORT,
+  SESSION_CLEANUP_SESSION,
+} session_cleanup_ntf_t;
+
 /*
  * Session states
  */