Adding pausing functionality for the profiler (#634)

* Dropping profiler session on stop request

* Changes to IXEventSession to simplify dropping sessions

* Stop sessions instead of dropping, disable flaky tests

* Initial framework for profiler pause requests

* Restructuring profiler session monitoring

* Fixes to session monitor

* Testing for pause functionality

* Fixing comments from PR

* Changes to testing

* Commenting out flaky test

* Deleting leftover testing code

src/Microsoft.SqlTools.ServiceLayer/Profiler/Contracts/PauseProfilingRequest.cs

@@ -0,0 +1,34 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Utility;
+using Microsoft.SqlTools.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.Profiler.Contracts
+{
+    /// <summary>
+    /// Pause Profiling request parameters
+    /// </summary>
+    public class PauseProfilingParams : GeneralRequestDetails
+    {
+        public string OwnerUri { get; set; }
+    }
+
+    public class PauseProfilingResult{}
+
+    /// <summary>
+    /// Pause Profile request type
+    /// </summary>
+    public class PauseProfilingRequest
+    {
+        /// <summary>
+        /// Request definition
+        /// </summary>
+        public static readonly
+            RequestType<PauseProfilingParams, PauseProfilingResult> Type =
+            RequestType<PauseProfilingParams, PauseProfilingResult>.Create("profiler/pause");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/Profiler/Contracts/StopProfilingRequest.cs

@@ -16,12 +16,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler.Contracts
         public string OwnerUri { get; set; }
     }
 
-    public class StopProfilingResult
+    public class StopProfilingResult{}
-    {
-        public bool Succeeded { get; set; }
-
-        public string ErrorMessage { get; set; }
-    }
 
     /// <summary>
     /// Start Profile request type

src/Microsoft.SqlTools.ServiceLayer/Profiler/IProfilerSessionMonitor.cs

@@ -22,11 +22,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         /// <summary>
         /// Starts monitoring a profiler session
         /// </summary>
-        bool StartMonitoringSession(ProfilerSession session);
+        bool StartMonitoringSession(string viewerId, IXEventSession session);
 
         /// <summary>
         /// Stops monitoring a profiler session
         /// </summary>
-        bool StopMonitoringSession(string sessionId, out ProfilerSession session);
+        bool StopMonitoringSession(string viewerId, out ProfilerSession session);
+
+        /// <summary>
+        /// Pauses or Unpauses the stream of events to the viewer
+        /// </summary>
+        void PauseViewer(string viewerId);
     }
 }

src/Microsoft.SqlTools.ServiceLayer/Profiler/IXEventSession.cs

@@ -11,13 +11,23 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
     public interface IXEventSession
     {
         /// <summary>
-        /// Reads XEvent XML from the default session target
+        /// Gets unique XEvent session Id
         /// </summary>
-        string GetTargetXml();
+        int Id { get; }
+
+        /// <summary>
+        /// Starts XEvent session
+        /// </summary>
+        void Start();
 
         /// <summary>
         /// Stops XEvent session
         /// </summary>
         void Stop();
+
+        /// <summary>
+        /// Reads XEvent XML from the default session target
+        /// </summary>
+        string GetTargetXml();
     }
 }

src/Microsoft.SqlTools.ServiceLayer/Profiler/IXEventSessionFactory.cs

@@ -16,8 +16,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
     public interface IXEventSessionFactory
     {
         /// <summary>
-        /// Create a new XEvent session
+        /// Gets or creates an XEvent session with the given template
         /// </summary>
-        IXEventSession CreateXEventSession(ConnectionInfo connInfo);
+        IXEventSession GetOrCreateXEventSession(string template, ConnectionInfo connInfo);
     }
 }

src/Microsoft.SqlTools.ServiceLayer/Profiler/ProfilerService.cs

@@ -110,6 +110,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
             this.ServiceHost = serviceHost;
             this.ServiceHost.SetRequestHandler(StartProfilingRequest.Type, HandleStartProfilingRequest);
             this.ServiceHost.SetRequestHandler(StopProfilingRequest.Type, HandleStopProfilingRequest);
+            this.ServiceHost.SetRequestHandler(PauseProfilingRequest.Type, HandlePauseProfilingRequest);
 
             this.SessionMonitor.AddSessionListener(this);
         }
@@ -129,8 +130,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
 
                 if (connInfo != null)
                 {
-                    ProfilerSession session = StartSession(parameters.OwnerUri, connInfo);
+                    int xEventSessionId = StartSession(parameters.OwnerUri, parameters.TemplateName, connInfo);
-                    result.SessionId = session.SessionId;
+                    result.SessionId = xEventSessionId.ToString();
                     result.Succeeded = true;
                 }
                 else
@@ -158,10 +159,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
                 monitor.StopMonitoringSession(parameters.OwnerUri, out session);
                 session.XEventSession.Stop();
 
-                await requestContext.SendResult(new StopProfilingResult
+                await requestContext.SendResult(new StopProfilingResult{});
-                {
-                    Succeeded = true
-                });
             }
             catch (Exception e)
             {
@@ -170,48 +168,53 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         }
 
         /// <summary>
-        /// Starts a new profiler session for the provided connection
+        /// Handle request to pause a profiling session
         /// </summary>
-        internal ProfilerSession StartSession(string sessionId, ConnectionInfo connInfo)
+        internal async Task HandlePauseProfilingRequest(PauseProfilingParams parameters, RequestContext<PauseProfilingResult> requestContext)
+        {
+            try
+            {
+                monitor.PauseViewer(parameters.OwnerUri);
+
+                await requestContext.SendResult(new PauseProfilingResult{});
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendError(e);
+            }
+        }
+
+        /// <summary>
+        /// Starts a new profiler session or connects to an existing session
+        /// for the provided connection and template info
+        /// </summary>
+        /// <returns>
+        /// The XEvent Session Id that was started
+        /// </returns>
+        internal int StartSession(string ownerUri, string template, ConnectionInfo connInfo)
         {
             // create a new XEvent session and Profiler session
-            var xeSession = this.XEventSessionFactory.CreateXEventSession(connInfo);
+            var xeSession = this.XEventSessionFactory.GetOrCreateXEventSession(template, connInfo);
-            var profilerSession = new ProfilerSession()
-            {
-                SessionId = sessionId,
-                XEventSession = xeSession
-            };
 
             // start monitoring the profiler session
-            monitor.StartMonitoringSession(profilerSession);
+            monitor.StartMonitoringSession(ownerUri, xeSession);
 
-            return profilerSession;
+            return xeSession.Id;
         }
 
         /// <summary>
-        /// Create a new XEvent sessions per the IXEventSessionFactory contract
+        /// Gets or creates an XEvent session with the given template per the IXEventSessionFactory contract
-        /// </summary>
-        public IXEventSession CreateXEventSession(ConnectionInfo connInfo)
-        {
-            var sqlConnection = ConnectionService.OpenSqlConnection(connInfo);
-            SqlStoreConnection connection = new SqlStoreConnection(sqlConnection);
-            Session session = ProfilerService.GetOrCreateSession(connection, "Profiler");
-
-            // create xevent session wrapper
-            return new XEventSession()
-            {
-                Session = session
-            };
-        }
-
-        /// <summary>
-        /// Gets an existing XEvent session or creates one if no matching session exists.
         /// Also starts the session if it isn't currently running
         /// </summary>
-        private static Session GetOrCreateSession(SqlStoreConnection connection, string sessionName)
+        public IXEventSession GetOrCreateXEventSession(string template, ConnectionInfo connInfo)
         {
+            string sessionName = "Profiler";
+
+            var sqlConnection = ConnectionService.OpenSqlConnection(connInfo);
+            SqlStoreConnection connection = new SqlStoreConnection(sqlConnection);
             XEStore store = new XEStore(connection);
             Session session = store.Sessions[sessionName];
+
             // start the session if it isn't already running
             if (session == null)
             {
@@ -222,7 +225,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
             {
                 session.Start();
             }
-            return session;
+
+            // create xevent session wrapper
+            return new XEventSession()
+            {
+                Session = session
+            };
         }
 
         private static Session CreateSession(SqlStoreConnection connection, string sessionName)

src/Microsoft.SqlTools.ServiceLayer/Profiler/ProfilerSession.cs

@@ -24,10 +24,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         private TimeSpan pollingDelay = DefaultPollingDelay;
         private ProfilerEvent lastSeenEvent = null;
 
-        /// <summary>
+        public bool pollImmediatly = false;
-        /// Unique ID for the session
-        /// </summary>
-        public string SessionId { get; set; }
 
         /// <summary>
         /// Connection to use for the session
@@ -47,10 +44,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         {
             lock (this.pollingLock)
             {
-                if (!this.isPolling && DateTime.Now.Subtract(this.lastPollTime) >= pollingDelay)
+                if (pollImmediatly || (!this.isPolling && DateTime.Now.Subtract(this.lastPollTime) >= pollingDelay))
                 {
                     this.isPolling = true;
                     this.lastPollTime = DateTime.Now;
+                    this.pollImmediatly = false;
                     return true;
                 }
                 else

src/Microsoft.SqlTools.ServiceLayer/Profiler/ProfilerSessionMonitor.cs

@@ -20,7 +20,7 @@ using Microsoft.SqlTools.Utility;
 namespace Microsoft.SqlTools.ServiceLayer.Profiler
 {
     /// <summary>
-    /// Classs to monitor active profiler sessions
+    /// Class to monitor active profiler sessions
     /// </summary>
     public class ProfilerSessionMonitor : IProfilerSessionMonitor
     {
@@ -30,9 +30,33 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
 
         private object listenersLock = new object();
 
+        private object pollingLock = new object();
+
         private Task processorThread = null;
 
-        private Dictionary<string, ProfilerSession> monitoredSessions = new Dictionary<string, ProfilerSession>();
+        private struct Viewer
+        {
+            public string Id { get; set; }
+            public bool active { get; set; }
+
+            public int xeSessionId { get; set; }
+
+            public Viewer(string Id, bool active, int xeId)
+            {
+                this.Id = Id;
+                this.active = active;
+                this.xeSessionId = xeId;
+            }
+        };
+
+        // XEvent Session Id's matched to the Profiler Id's watching them
+        private Dictionary<int, List<string>> sessionViewers = new Dictionary<int, List<string>>();
+
+        // XEvent Session Id's matched to their Profiler Sessions
+        private Dictionary<int, ProfilerSession> monitoredSessions = new Dictionary<int, ProfilerSession>();
+
+        // ViewerId -> Viewer objects
+        private Dictionary<string, Viewer> allViewers = new Dictionary<string, Viewer>();
 
         private List<IProfilerSessionListener> listeners = new List<IProfilerSessionListener>();
 
@@ -48,21 +72,50 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         }
 
         /// <summary>
-        /// Start monitoring the provided sessions
+        /// Start monitoring the provided session
         /// </summary>
-        public bool StartMonitoringSession(ProfilerSession session)
+        public bool StartMonitoringSession(string viewerId, IXEventSession session)
         {
             lock (this.sessionsLock)
             {
                 // start the monitoring thread
                 if (this.processorThread == null)
                 {
-                    this.processorThread = Task.Factory.StartNew(ProcessSessions);;
+                    this.processorThread = Task.Factory.StartNew(ProcessSessions);
                 }
 
-                if (!this.monitoredSessions.ContainsKey(session.SessionId))
+                // create new profiling session if needed
+                if (!this.monitoredSessions.ContainsKey(session.Id))
                 {
-                    this.monitoredSessions.Add(session.SessionId, session);
+                    var profilerSession = new ProfilerSession();
+                    profilerSession.XEventSession = session;
+
+                    this.monitoredSessions.Add(session.Id, profilerSession);
+                }
+
+                // create a new viewer, or configure existing viewer
+                Viewer viewer;
+                if (!this.allViewers.TryGetValue(viewerId, out viewer))
+                {
+                    viewer = new Viewer(viewerId, true, session.Id);
+                    allViewers.Add(viewerId, viewer);
+                }
+                else
+                {
+                    viewer.active = true;
+                    viewer.xeSessionId = session.Id;
+                }
+
+                // add viewer to XEvent session viewers
+                List<string> viewers;
+                if (this.sessionViewers.TryGetValue(session.Id, out viewers))
+                {
+                    viewers.Add(viewerId);
+                }
+                else
+                {
+                    viewers = new List<string>{ viewerId };
+                    sessionViewers.Add(session.Id, viewers);
                 }
             }
 
@@ -70,15 +123,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         }
 
         /// <summary>
-        /// Stop monitoring the session specified by the sessionId
+        /// Stop monitoring the session watched by viewerId
         /// </summary>
-        public bool StopMonitoringSession(string sessionId, out ProfilerSession session)
+        public bool StopMonitoringSession(string viewerId, out ProfilerSession session)
         {
             lock (this.sessionsLock)
             {
-                if (this.monitoredSessions.ContainsKey(sessionId))
+                Viewer v;
+                if (this.allViewers.TryGetValue(viewerId, out v))
                 {
-                    return this.monitoredSessions.Remove(sessionId, out session);
+                    return RemoveSession(v.xeSessionId, out session);
                 }
                 else
                 {
@@ -88,6 +142,61 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
             }
         }
 
+        /// <summary>
+        /// Toggle the pause state for the viewer
+        /// </summary>
+        public void PauseViewer(string viewerId)
+        {
+            lock (this.sessionsLock)
+            {
+                Viewer v = this.allViewers[viewerId];
+                v.active = !v.active;
+                this.allViewers[viewerId] = v;
+            }
+        }
+
+        private bool RemoveSession(int sessionId, out ProfilerSession session)
+        {
+            lock (this.sessionsLock)
+            {
+                if (this.monitoredSessions.Remove(sessionId, out session))
+                {
+                    //remove all viewers for this session
+                    List<string> viewerIds;
+                    if (sessionViewers.Remove(sessionId, out viewerIds))
+                    {
+                        foreach (String viewerId in viewerIds)
+                        {
+                            this.allViewers.Remove(viewerId);
+                        }
+                        return true;
+                    }
+                    else
+                    {
+                        session = null;
+                        return false;
+                    }
+                }
+                else
+                {
+                    session = null;
+                    return false;
+                }
+            }
+        }
+
+        public void PollSession(int sessionId)
+        {
+            lock (this.sessionsLock)
+            {
+                this.monitoredSessions[sessionId].pollImmediatly = true;
+            }
+            lock (this.pollingLock)
+            {
+                Monitor.Pulse(pollingLock);
+            }
+        }
+
         /// <summary>
         /// The core queue processing method
         /// </summary>
@@ -96,15 +205,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
         {
             while (true)
             {
-                lock (this.sessionsLock)
+                lock (this.pollingLock)
                 {
-                    foreach (var session in this.monitoredSessions.Values)
+                    lock (this.sessionsLock)
                     {
-                        ProcessSession(session);
+                        foreach (var session in this.monitoredSessions.Values)
+                        {
+                            List<string> viewers = this.sessionViewers[session.XEventSession.Id];
+                            if (viewers.Any(v => allViewers[v].active))
+                            {
+                                ProcessSession(session);
+                            }
+                        }
                     }
+                    Monitor.Wait(this.pollingLock, PollingLoopDelay);
                 }
-
-                Thread.Sleep(PollingLoopDelay);
             }
         }
 
@@ -120,7 +235,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
                     var events = PollSession(session);
                     if (events.Count > 0)
                     {
-                        SendEventsToListeners(session.SessionId, events);
+                        // notify all viewers for the polled session
+                        List<string> viewerIds = this.sessionViewers[session.XEventSession.Id];
+                        foreach (string viewerId in viewerIds)
+                        {
+                            if (allViewers[viewerId].active)
+                            {
+                                SendEventsToListeners(viewerId, events);
+                            }
+                        }
                     }
                 });
             }
@@ -151,7 +274,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Profiler
                     }
                 }
             }
-            catch(Exception ex)
+            catch (Exception ex)
             {
                 Logger.Write(LogLevel.Warning, "Failed to pool session. error: " + ex.Message);
             }

src/Microsoft.SqlTools.ServiceLayer/Profiler/XEventSession.cs

@@ -10,10 +10,23 @@ using Microsoft.SqlTools.ServiceLayer.Profiler.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.Profiler
 {
+    /// <summary>
+    /// Class to access underlying XEvent session.
+    /// </summary>
     public class XEventSession : IXEventSession
     {
         public Session Session { get; set; }
 
+        public int Id
+        {
+            get { return Session.ID; }
+        }
+
+        public void Start()
+        {
+            this.Session.Start();
+        }
+
         public void Stop()
         {
             this.Session.Stop();

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Profiler/ProfilerServiceTests.cs

@@ -82,7 +82,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Profiler
         {
             var liveConnection = LiveConnectionHelper.InitLiveConnectionInfo("master");
             ProfilerService profilerService = new ProfilerService();
-            IXEventSession xeSession = profilerService.CreateXEventSession(liveConnection.ConnectionInfo);
+            IXEventSession xeSession = profilerService.GetOrCreateXEventSession("Profiler", liveConnection.ConnectionInfo);
             Assert.NotNull(xeSession);
             Assert.NotNull(xeSession.GetTargetXml());
         }

test/Microsoft.SqlTools.ServiceLayer.UnitTests/Profiler/ProfilerServiceTests.cs

@@ -7,6 +7,7 @@ using System;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Collections.Generic;
 using Microsoft.SqlServer.Management.XEvent;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Connection;
@@ -32,6 +33,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
         public async Task TestStartProfilingRequest()
         {
             string sessionId = null;
+            bool recievedEvents = false;
             string testUri = "profiler_uri";
             var requestContext = new Mock<RequestContext<StartProfilingResult>>();
             requestContext.Setup(rc => rc.SendResult(It.IsAny<StartProfilingResult>()))
@@ -42,10 +44,15 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
                     return Task.FromResult(0);
                 });
 
-            var sessionListener = new TestSessionListener();
+            // capture listener event notifications
+            var mockListener = new Mock<IProfilerSessionListener>();
+            mockListener.Setup(p => p.EventsAvailable(It.IsAny<string>(), It.IsAny<List<ProfilerEvent>>())).Callback(() =>
+                {
+                    recievedEvents = true;
+                });
 
             var profilerService = new ProfilerService();
-            profilerService.SessionMonitor.AddSessionListener(sessionListener);
+            profilerService.SessionMonitor.AddSessionListener(mockListener.Object);
             profilerService.ConnectionServiceInstance = TestObjects.GetTestConnectionService();
             ConnectionInfo connectionInfo = TestObjects.GetTestConnectionInfo();
             profilerService.ConnectionServiceInstance.OwnerToConnectionMap.Add(testUri, connectionInfo);
@@ -55,15 +62,33 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
             requestParams.OwnerUri = testUri;
             requestParams.TemplateName = "Standard";
 
+            // start profiling session
             await profilerService.HandleStartProfilingRequest(requestParams, requestContext.Object);
 
-            // wait a bit for profile sessions to be polled
+            profilerService.SessionMonitor.PollSession(1);
-            Thread.Sleep(500);
+            // simulate a short polling delay
+            Thread.Sleep(200);
+            profilerService.SessionMonitor.PollSession(1);
+
+            // wait for polling to finish, or for timeout
+            System.Timers.Timer pollingTimer = new System.Timers.Timer();
+            pollingTimer.Interval = 10000;
+            pollingTimer.Start();
+            bool timeout = false;
+            pollingTimer.Elapsed += new System.Timers.ElapsedEventHandler((s_, e_) => {timeout = true;});
+            while (sessionId == null && !timeout)
+            {
+                Thread.Sleep(250);
+            }
+            pollingTimer.Stop();
 
             requestContext.VerifyAll();
 
-            Assert.Equal(sessionListener.PreviousSessionId, sessionId);
+            // Check that the correct XEvent session was started
-            Assert.Equal(sessionListener.PreviousEvents.Count, 1);
+            Assert.Equal(sessionId, "1");
+
+            // check that the proper owner Uri was used
+            Assert.True(recievedEvents);
         }
 
         /// <summary>
@@ -82,11 +107,11 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
             requestContext.Setup(rc => rc.SendResult(It.IsAny<StopProfilingResult>()))
                 .Returns<StopProfilingResult>((result) =>
                 {
-                    success = result.Succeeded;
+                    success = true;
                     return Task.FromResult(0);
                 });
 
-            // capture if session was dropped
+            // capture if session was stopped
             var mockSession = new Mock<IXEventSession>();
             mockSession.Setup(p => p.Stop()).Callback(() =>
                 {
@@ -104,17 +129,13 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
             var requestParams = new StopProfilingParams();
             requestParams.OwnerUri = testUri;
 
-            ProfilerSession session = new ProfilerSession();
+            profilerService.SessionMonitor.StartMonitoringSession(testUri, mockSession.Object);
-            session.XEventSession = mockSession.Object;
-            session.SessionId = testUri;
-
-            profilerService.SessionMonitor.StartMonitoringSession(session);
 
             await profilerService.HandleStopProfilingRequest(requestParams, requestContext.Object);
 
             requestContext.VerifyAll();
 
-            // check that session was succesfully stopped and drop was called
+            // check that session was succesfully stopped and stop was called
             Assert.True(success);
             Assert.True(stopped);
 
@@ -122,5 +143,97 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
             ProfilerSession ps;
             Assert.False(profilerService.SessionMonitor.StopMonitoringSession(testUri, out ps));
         }
+
+        /// <summary>
+        /// Test pausing then resuming a session
+        /// </summary>
+        /// <returns></returns>
+        [Fact]
+        public async Task TestPauseProfilingRequest()
+        {
+            bool success = false;
+            string testUri = "test_session";
+            bool recievedEvents = false;
+
+            // capture pausing results
+            var requestContext = new Mock<RequestContext<PauseProfilingResult>>();
+            requestContext.Setup(rc => rc.SendResult(It.IsAny<PauseProfilingResult>()))
+                .Returns<PauseProfilingResult>((result) =>
+                {
+                    success = true;
+                    return Task.FromResult(0);
+                });
+
+            // capture listener event notifications
+            var mockListener = new Mock<IProfilerSessionListener>();
+            mockListener.Setup(p => p.EventsAvailable(It.IsAny<string>(), It.IsAny<List<ProfilerEvent>>())).Callback(() =>
+                {
+                    recievedEvents = true;
+                });
+
+            // setup profiler service
+            var profilerService = new ProfilerService();
+            profilerService.SessionMonitor.AddSessionListener(mockListener.Object);
+            profilerService.ConnectionServiceInstance = TestObjects.GetTestConnectionService();
+            ConnectionInfo connectionInfo = TestObjects.GetTestConnectionInfo();
+            profilerService.ConnectionServiceInstance.OwnerToConnectionMap.Add(testUri, connectionInfo);
+
+            var requestParams = new PauseProfilingParams();
+            requestParams.OwnerUri = testUri;
+
+            // begin monitoring session
+            profilerService.SessionMonitor.StartMonitoringSession(testUri, new TestXEventSession1());
+
+            // poll the session
+            profilerService.SessionMonitor.PollSession(1);
+            Thread.Sleep(500);
+            profilerService.SessionMonitor.PollSession(1);
+
+            // wait for polling to finish, or for timeout
+            System.Timers.Timer pollingTimer = new System.Timers.Timer();
+            pollingTimer.Interval = 10000;
+            pollingTimer.Start();
+            bool timeout = false;
+            pollingTimer.Elapsed += new System.Timers.ElapsedEventHandler((s_, e_) => {timeout = true;});
+            while (!recievedEvents && !timeout)
+            {
+                Thread.Sleep(250);
+            }
+            pollingTimer.Stop();
+
+            // confirm that polling works
+            Assert.True(recievedEvents);
+
+            // pause viewer
+            await profilerService.HandlePauseProfilingRequest(requestParams, requestContext.Object);
+            Assert.True(success);
+
+            recievedEvents = false;
+            success = false;
+
+            profilerService.SessionMonitor.PollSession(1);
+
+            // confirm that no events were sent to paused listener
+            Assert.False(recievedEvents);
+
+            // unpause viewer
+            await profilerService.HandlePauseProfilingRequest(requestParams, requestContext.Object);
+            Assert.True(success);
+
+            profilerService.SessionMonitor.PollSession(1);
+
+            // wait for polling to finish, or for timeout
+            timeout = false;
+            pollingTimer.Start();
+            while (!recievedEvents && !timeout)
+            {
+                Thread.Sleep(250);
+            }
+
+            // check that events got sent to listener
+            Assert.True(recievedEvents);
+
+            requestContext.VerifyAll();
+        }
     }
 }

test/Microsoft.SqlTools.ServiceLayer.UnitTests/Profiler/ProfilerTestObjects.cs

@@ -188,19 +188,168 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Profiler
             "	</event>" +
             "</RingBufferTarget>";
 
+
+
+        public int Id { get { return 51; } }
+
+        public void Start(){}
+
+        public void Stop(){}
         public string GetTargetXml()
         {
             return testXEventXml;
         }
+    }
+
+    public class TestXEventSession1 : IXEventSession
+    {
+        private const string testXEventXml_1 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        private const string testXEventXml_2 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-10-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        private const string testXEventXml_3 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-10-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-11-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>1</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        public int Id { get { return 1; } }
+
+        public void Start(){}
 
         public void Stop(){}
+
+        private int pollCount = 0;
+        private string[] poll_returns = { testXEventXml_1, testXEventXml_2, testXEventXml_3 };
+        public string GetTargetXml()
+        {
+            string res = poll_returns[pollCount];
+            pollCount++;
+            pollCount = pollCount > 2 ? 0 : pollCount;
+            return res;
+        }
+    }
+
+        public class TestXEventSession2 : IXEventSession
+    {
+        private const string testXEventXml_1 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        private const string testXEventXml_2 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-10-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        private const string testXEventXml_3 =
+            "<RingBufferTarget truncated=\"0\" processingTime=\"3\" totalEventsProcessed=\"1\" eventCount=\"1\" droppedCount=\"0\" memoryUsed=\"47996\">" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-09-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-10-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "	<event name=\"existing_connection\" package=\"sqlserver\" timestamp=\"2017-11-08T07:46:53.579Z\">" +
+            "		<data name=\"session_id\">" +
+            "			<type name=\"int16\" package=\"package0\"></type>" +
+            "			<value>2</value>" +
+            "		</data>" +
+            "	</event>" +
+            "</RingBufferTarget>";
+
+        public int Id { get { return 2; } }
+
+        public void Start(){}
+
+        public void Stop(){}
+
+        private int pollCount = 0;
+        private string[] poll_returns = { testXEventXml_1, testXEventXml_2, testXEventXml_3 };
+        public string GetTargetXml()
+        {
+            string res = poll_returns[pollCount];
+            pollCount++;
+            pollCount = pollCount > 2 ? 0 : pollCount;
+            return res;
+        }
     }
 
     public class TestXEventSessionFactory : IXEventSessionFactory
     {
-        public IXEventSession CreateXEventSession(ConnectionInfo connInfo)
+        private int sessionNum = 1;
+        public IXEventSession GetOrCreateXEventSession(string template, ConnectionInfo connInfo)
         {
-            return new TestXEventSession();
+            if(sessionNum == 1)
+            {
+                sessionNum = 2;
+                return new TestXEventSession1();
+            }
+            else
+            {
+                sessionNum = 1;
+                return new TestXEventSession2();
+            }
         }
     }
 }