From 19aff26259edab2be636c2d0ea52a9e3b4f82324 Mon Sep 17 00:00:00 2001
From: Egor Matveev <e.matveev@hh.ru>
Date: Wed, 1 Jan 2025 13:53:05 +0300
Subject: [PATCH 1/2] fix

---
 main.go | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index 8d7344a..8fabc37 100644
--- a/main.go
+++ b/main.go
@@ -1,15 +1,26 @@
 package main
 
 import (
+	"log"
 	"net/http"
 	"queues-go/app/routers"
 	client "queues-go/app/storage/mongo"
+	"time"
 )
 
+func handlerWrapper(f func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		start := time.Now()
+		f(w, r)
+		log.Printf("%s %s", r.URL, time.Since(start))
+	}
+}
+
 func main() {
 	client.Connect()
-	http.HandleFunc("/api/v1/take", routers.Take)
-	http.HandleFunc("/api/v1/finish", routers.Finish)
-	http.HandleFunc("/api/v1/put", routers.Put)
+	http.HandleFunc("/api/v1/take", handlerWrapper(routers.Take))
+	http.HandleFunc("/api/v1/finish", handlerWrapper(routers.Finish))
+	http.HandleFunc("/api/v1/put", handlerWrapper(routers.Put))
+	log.Printf("Server started")
 	http.ListenAndServe(":1239", nil)
 }

From 3504e132addbab430db2690c2d7d7805ca27a1b8 Mon Sep 17 00:00:00 2001
From: Egor Matveev <e.matveev@hh.ru>
Date: Wed, 1 Jan 2025 14:18:19 +0300
Subject: [PATCH 2/2] fix

---
 app/routers/finish.go | 10 ++++------
 app/routers/put.go    | 10 ++++------
 app/routers/take.go   | 16 +++-------------
 main.go               | 18 +++++++++++++++---
 4 files changed, 26 insertions(+), 28 deletions(-)

diff --git a/app/routers/finish.go b/app/routers/finish.go
index b471b8b..5cc1ad6 100644
--- a/app/routers/finish.go
+++ b/app/routers/finish.go
@@ -10,20 +10,18 @@ type FinishRequestBody struct {
 	Id string `json:"id"`
 }
 
-func Finish(w http.ResponseWriter, r *http.Request) {
+func Finish(r *http.Request) (interface{}, int) {
 	d := json.NewDecoder(r.Body)
 	body := FinishRequestBody{}
 	err := d.Decode(&body)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		return
+		return nil, http.StatusBadRequest
 	}
 
 	err = tasks.Finish(body.Id)
 	if err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
-	w.WriteHeader(http.StatusAccepted)
+	return nil, http.StatusAccepted
 }
diff --git a/app/routers/put.go b/app/routers/put.go
index cfd9834..75f2cc1 100644
--- a/app/routers/put.go
+++ b/app/routers/put.go
@@ -15,13 +15,12 @@ type PutRequestBody struct {
 	Delay            *int   `json:"delay"`
 }
 
-func Put(w http.ResponseWriter, r *http.Request) {
+func Put(r *http.Request) (interface{}, int) {
 	d := json.NewDecoder(r.Body)
 	body := PutRequestBody{}
 	err := d.Decode(&body)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		return
+		return nil, http.StatusBadRequest
 	}
 
 	queue := r.Header.Get("queue")
@@ -44,9 +43,8 @@ func Put(w http.ResponseWriter, r *http.Request) {
 
 	err = tasks.Add(task)
 	if err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
-	w.WriteHeader(http.StatusAccepted)
+	return nil, http.StatusAccepted
 }
diff --git a/app/routers/take.go b/app/routers/take.go
index 90e4545..42bc6ff 100644
--- a/app/routers/take.go
+++ b/app/routers/take.go
@@ -1,7 +1,6 @@
 package routers
 
 import (
-	"encoding/json"
 	"net/http"
 	tasks "queues-go/app/storage/mongo/collections"
 )
@@ -16,13 +15,11 @@ type TakeResponse struct {
 	Task *TaskResponse `json:"task"`
 }
 
-func Take(w http.ResponseWriter, r *http.Request) {
+func Take(r *http.Request) (interface{}, int) {
 	queue := r.Header.Get("queue")
 	task, err := tasks.Take(queue)
 	if err != nil {
-		println("Error taking")
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
 	var response TakeResponse
@@ -36,12 +33,5 @@ func Take(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	data, err := json.Marshal(response)
-	if err != nil {
-		println("Error marshalling")
-		w.WriteHeader(http.StatusInternalServerError)
-		return
-	}
-
-	w.Write(data)
+	return response, http.StatusOK
 }
diff --git a/main.go b/main.go
index 8fabc37..d2562b5 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"log"
 	"net/http"
 	"queues-go/app/routers"
@@ -8,11 +9,22 @@ import (
 	"time"
 )
 
-func handlerWrapper(f func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
+func handlerWrapper(f func(*http.Request) (interface{}, int)) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 		start := time.Now()
-		f(w, r)
-		log.Printf("%s %s", r.URL, time.Since(start))
+		response, status := f(r)
+		if response != nil {
+			data, err := json.Marshal(response)
+			if err != nil {
+				log.Printf("Error marshalling response")
+				status = http.StatusInternalServerError
+			}
+			w.WriteHeader(status)
+			w.Write(data)
+		} else {
+			w.WriteHeader(status)
+		}
+		log.Printf("%s %d %s", r.URL, status, time.Since(start))
 	}
 }