From 2ab1d9d774265c050383382fbc6187d6b2789ad6 Mon Sep 17 00:00:00 2001
From: Alex X <alexey.khit@gmail.com>
Date: Wed, 29 May 2024 17:32:11 +0300
Subject: [PATCH] Add handling if mp4 client drops connection

---
 internal/mp4/mp4.go | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/internal/mp4/mp4.go b/internal/mp4/mp4.go
index 78708a35f..2f59ba041 100644
--- a/internal/mp4/mp4.go
+++ b/internal/mp4/mp4.go
@@ -1,6 +1,7 @@
 package mp4
 
 import (
+	"context"
 	"net/http"
 	"strconv"
 	"strings"
@@ -127,20 +128,20 @@ func handlerMP4(w http.ResponseWriter, r *http.Request) {
 		header.Set("Content-Disposition", `attachment; filename="`+filename+`"`)
 	}
 
-	var duration *time.Timer
-	if s := query.Get("duration"); s != "" {
-		if i, _ := strconv.Atoi(s); i > 0 {
-			duration = time.AfterFunc(time.Second*time.Duration(i), func() {
-				_ = cons.Stop()
-			})
-		}
+	ctx := r.Context() // handle when the client drops the connection
+
+	if i := core.Atoi(query.Get("duration")); i > 0 {
+		timeout := time.Second * time.Duration(i)
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
 	}
 
-	_, _ = cons.WriteTo(w)
+	go func() {
+		<-ctx.Done()
+		_ = cons.Stop()
+		stream.RemoveConsumer(cons)
+	}()
 
-	stream.RemoveConsumer(cons)
-
-	if duration != nil {
-		duration.Stop()
-	}
+	_, _ = cons.WriteTo(w)
 }