improving cleanup in handlers

app/handlers.go

@@ -18,7 +18,7 @@ import (
 // GetHandlersChain returns the handlers chaining
 // thanks to the alice library.
 func GetHandlersChain() http.Handler {
-	return alice.New(enforceContentLengthHandler, enforceContentTypeHandler, convertHandler, serveHandler).ThenFunc(clearHandler)
+	return alice.New(enforceContentLengthHandler, enforceContentTypeHandler, convertHandler).ThenFunc(serveHandler)
 }
 
 type requestHasNoContentError struct{}
@@ -62,6 +62,7 @@ func convertHandler(next http.Handler) http.Handler {
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			logger.Error(err)
+			cleanup(r)
 			return
 		}
 
@@ -69,6 +70,7 @@ func convertHandler(next http.Handler) http.Handler {
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			logger.Error(err)
+			cleanup(r)
 			return
 		}
 
@@ -80,41 +82,54 @@ func convertHandler(next http.Handler) http.Handler {
 }
 
 // serveHandler simply serves the created PDF.
-func serveHandler(next http.Handler) http.Handler {
+func serveHandler(w http.ResponseWriter, r *http.Request) {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	path, err := context.GetResultFilePath(r)
-		path, err := context.GetResultFilePath(r)
+	if err != nil {
-		if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+		logger.Error(err)
-			logger.Error(err)
+		cleanup(r)
-		}
+		return
+	}
 
-		reader, err := os.Open(path)
+	reader, err := os.Open(path)
-		if err != nil {
+	if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+		http.Error(w, err.Error(), http.StatusInternalServerError)
-			logger.Error(err)
+		logger.Error(err)
-			return
+		cleanup(r)
-		}
+		return
+	}
 
-		defer reader.Close()
+	defer reader.Close()
 
-		resultFileInfo, err := reader.Stat()
+	resultFileInfo, err := reader.Stat()
-		if err != nil {
+	if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+		http.Error(w, err.Error(), http.StatusInternalServerError)
-			logger.Error(err)
+		logger.Error(err)
-			return
+		cleanup(r)
-		}
+		return
+	}
 
+	done := make(chan error, 1)
+	go func() {
 		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", resultFileInfo.Name()))
 		w.Header().Set("Content-Type", "application/pdf")
 		w.Header().Set("Content-Length", fmt.Sprintf("%d", resultFileInfo.Size()))
-		io.Copy(w, reader)
+		_, err := io.Copy(w, reader)
 
-		next.ServeHTTP(w, r)
+		done <- err
-	})
+	}()
+
+	err = <-done
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		logger.Error(err)
+	}
+
+	cleanup(r)
 }
 
-// clearHandler removes all files created during the conversion.
+// cleanup removes all files created during the conversion.
-func clearHandler(w http.ResponseWriter, r *http.Request) {
+func cleanup(r *http.Request) {
 	c, err := context.GetConverter(r)
 	if err != nil {
 		logger.Warn(err.Error())

app/handlers_test.go

@@ -11,7 +11,6 @@ import (
 
 	"github.com/thecodingmachine/gotenberg/app/config"
 	"github.com/thecodingmachine/gotenberg/app/context"
-	"github.com/thecodingmachine/gotenberg/app/converter"
 	"github.com/thecodingmachine/gotenberg/app/converter/process"
 	ghttp "github.com/thecodingmachine/gotenberg/app/http"
 
@@ -133,7 +132,7 @@ func TestConvertHandler(t *testing.T) {
 }
 
 func TestServeHandler(t *testing.T) {
-	h := alice.New(serveHandler).ThenFunc(fakeSuccessHandler)
+	h := alice.New().ThenFunc(serveHandler)
 
 	// case 1: sends a request without a result file path entry in its context.
 	req := httptest.NewRequest(http.MethodPost, "/", nil)
@@ -160,15 +159,3 @@ func TestServeHandler(t *testing.T) {
 		t.Errorf("Handler returned a wrong status code: got %v want %v", status, http.StatusOK)
 	}
 }
-
-func TestClearHandler(t *testing.T) {
-	// case 1: sends a request without a converter entry in its context.
-	req := httptest.NewRequest(http.MethodPost, "/", nil)
-	rr := httptest.NewRecorder()
-	clearHandler(rr, req)
-
-	// case 2: sends with a wrong converter entry in its context.
-	req = context.WithConverter(httptest.NewRequest(http.MethodPost, "/", nil), &converter.Converter{})
-	rr = httptest.NewRecorder()
-	clearHandler(rr, req)
-}