Use context object for socket close instead of notifier

streaming/connection.go

@@ -80,7 +80,9 @@ func (conn *Connection) Close() error {
 }
 
 // Serve starts serving data to a client, continuously feeding packets from the queue.
-func (conn *Connection) Serve() {
+// The context argument is used to react to Done() status - you should pass the request's Context() object here.
+// Returns true if we exited because the queue was closed, false if the external connection was closed instead.
+func (conn *Connection) Serve(ctx context.Context) bool {
 	// set the content type (important)
 	conn.writer.Header().Set("Content-Type", "video/mpeg")
 	// a stream is always current
@@ -108,16 +110,6 @@ func (conn *Connection) Serve() {
 		"message", "Sent header",
 	)
 
-	// see if can get notified about connection closure
-	notifier, ok := conn.writer.(http.CloseNotifier)
-	if !ok {
-		logger.Logkv(
-			"event", eventConnectionError,
-			"error", errorConnectionNoCloseNotify,
-			"message", "Writer does not support CloseNotify",
-		)
-	}
-
 	// this is the exit status indicator
 	qclosed := false
 	// start reading packets
@@ -151,7 +143,7 @@ func (conn *Connection) Serve() {
 				// indicate that the queue was closed
 				qclosed = true
 			}
-		case <-notifier.CloseNotify():
+		case <-ctx.Done():
 			// connection closed while we were waiting for more data
 			logger.Logkv(
 				"event", eventConnectionClosedWait,

streaming/streamer.go

@@ -345,6 +345,9 @@ func (streamer *Streamer) Stream(queue <-chan protocol.MpegTsPacket) error {
 				// close all downstream connections
 				for conn := range pool {
 					conn.Close()
+					// avoid waiting for the removal round-trip, this will make us less racy
+					// double deletes are safe, so nothing bad will happen when we do get the remove command later
+					delete(pool, request.Connection)
 				}
 				// TODO implement inhibit in the check api
 			case StreamerCommandAllow:
@@ -434,8 +437,9 @@ func (streamer *Streamer) ServeHTTP(writer http.ResponseWriter, request *http.Re
 			"remote", request.RemoteAddr,
 		)
 
+		// here's where the action happens
 		start := time.Now()
-		conn.Serve()
+		conn.Serve(request.Context())
 		duration := time.Since(start)
 
 		// done, remove the stale connection