Event sending thread logic and shutdown behaviour

Author: AndrewBurian

client.go

@@ -3,16 +3,22 @@ package eventsource
 import (
 	"io"
 	"net/http"
+	"sync"
 )
 
 type Client struct {
-	flush http.Flusher
-	write io.Writer
+	flush  http.Flusher
+	write  io.Writer
+	close  http.CloseNotifier
+	events chan *Event
+	closed bool
+	waiter sync.WaitGroup
 }
 
 func NewClient(w http.ResponseWriter) *Client {
 	c := &Client{
-		write: w,
+		events: make(chan *Event, 1),
+		write:  w,
 	}
 
 	// Check to ensure we support flushing
@@ -22,14 +28,68 @@ func NewClient(w http.ResponseWriter) *Client {
 	}
 	c.flush = flush
 
+	// Check to ensure we support close notifications
+	closer, ok := w.(http.CloseNotifier)
+	if !ok {
+		return nil
+	}
+	c.close = closer
+
+	// start the sending thread
+	c.waiter.Add(1)
+	go c.run()
 	return c
 }
 
-func (c *Client) Send(ev *Event) {
-	io.Copy(c.write, ev)
-	c.flush.Flush()
+// Send queues an event to be sent to the client.
+// This does not block until the event has been sent.
+// Returns an error if the Client has disconnected
+func (c *Client) Send(ev *Event) error {
+	if c.closed {
+		return io.ErrClosedPipe
+	}
+	c.events <- ev
+	return nil
+}
+
+// Shutdown terminates a client connection
+func (c *Client) Shutdown() {
+	close(c.events)
+	c.waiter.Wait()
 }
 
+// Wait blocks and waits for the client to be shutdown.
+// Call this is http handler threads to prevent the server from closing
+// the client connection.
+func (c *Client) Wait() {
+	c.waiter.Wait()
+}
+
+// Worker thread for the client responsible for writing events
 func (c *Client) run() {
 
+	for {
+		select {
+		case ev, ok := <-c.events:
+			// check for shutdown
+			if !ok {
+				c.closed = true
+				c.waiter.Done()
+				return
+			}
+
+			// send the event
+			io.Copy(c.write, ev)
+			c.flush.Flush()
+
+		case _ = <-c.close.CloseNotify():
+			c.closed = true
+			// empty the event channel
+			for _, ok := <-c.events; ok != false; {
+			}
+			c.waiter.Done()
+			return
+		}
+
+	}
 }

event.go

@@ -17,7 +17,7 @@ type Event struct {
 	bufSet bool
 }
 
-// DataEvent creates a new Event with the data set
+// DataEvent creates a new Event with the data field set
 func DataEvent(data string) *Event {
 	e := &Event{}
 	io.WriteString(e, data)

stream.go

@@ -9,7 +9,6 @@ import (
 type Stream struct {
 	clients           list.List
 	listLock          sync.RWMutex
-	broadcast         chan *Event
 	shutdownWait      sync.WaitGroup
 	clientConnectHook func(*http.Request, *Client)
 }
@@ -21,7 +20,6 @@ type registeredClient struct {
 
 func New() *Stream {
 	s := &Stream{}
-	go s.run()
 	return s
 }
 
@@ -82,10 +80,17 @@ func (s *Stream) Publish(topic string, e *Event) {
 	}
 }
 
+// Shutdown terminates all clients connected to the stream and removes them
 func (s *Stream) Shutdown() {
-
+	for element := s.clients.Front(); element != nil; element.Next() {
+		cli := element.Value.(*registeredClient)
+		cli.c.Shutdown()
+		s.clients.Remove(element)
+	}
 }
 
+// CloseTopic removes all client associations with this topic, but does not
+// terminate them or remove
 func (s *Stream) CloseTopic(topic string) {
 
 }
@@ -104,27 +109,6 @@ func (s *Stream) ClientConnectHook(fn func(*http.Request, *Client)) {
 	s.clientConnectHook = fn
 }
 
-func (s *Stream) run() {
-
-	for {
-		select {
-
-		case ev, ok := <-s.broadcast:
-
-			// end of the broadcast channel indicates a shutdown
-			if !ok {
-				//s.closeAll()
-				s.shutdownWait.Done()
-				return
-			}
-
-			// otherwise normal message
-			s.sendAll(ev)
-
-		}
-	}
-}
-
 func (s *Stream) sendAll(ev *Event) {
 
 }