From 37d5e4bb8094f87fb30ae8f77345956b7d729e2d Mon Sep 17 00:00:00 2001
From: 4rkal <4rkal@horsefucker.org>
Date: Mon, 23 Dec 2024 16:39:33 +0200
Subject: [PATCH] minor fixed for postgres

---
 cmd/main.go        | 17 +++++++++--------
 docker-compose.yml |  3 ++-
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index f43f25c..2d99de6 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -11,6 +11,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/4rkal/shortr/models"
@@ -32,6 +33,7 @@ type StatsFormData struct {
 }
 
 var linkMap = map[string]*models.Link{}
+var linkMapMutex sync.RWMutex
 
 var baseurl *string
 var db *sql.DB
@@ -46,10 +48,6 @@ func init() {
 		panic(fmt.Sprintf("Failed to connect to database: %v", err))
 	}
 
-	if err := db.Ping(); err != nil {
-		panic(fmt.Sprintf("Failed to ping database: %v", err))
-	}
-
 	if err := loadLinksIntoMemory(); err != nil {
 		panic(fmt.Sprintf("Failed to load links into memory: %v", err))
 	}
@@ -64,10 +62,6 @@ func init() {
 	if err != nil {
 		log.Fatal("Error creating table: ", err)
 	}
-
-	if err := loadLinksIntoMemory(); err != nil {
-		log.Fatalf("Failed to load links into memory: %v", err)
-	}
 }
 
 func main() {
@@ -90,7 +84,9 @@ func main() {
 func RedirectHandler(c echo.Context) error {
 	id := c.Param("id")
 
+	linkMapMutex.RLock()
 	link, found := linkMap[id]
+	linkMapMutex.RUnlock()
 	if !found {
 		return c.String(http.StatusNotFound, "Link not found")
 	}
@@ -134,7 +130,9 @@ func SubmitHandler(c echo.Context) error {
 		}
 	}
 
+	linkMapMutex.Lock()
 	linkMap[id] = &models.Link{Id: id, Url: data.Url}
+	linkMapMutex.Unlock()
 
 	_, err := db.Exec("INSERT INTO links (id, url, clicks) VALUES ($1, $2, $3)", id, data.Url, 0)
 	if err != nil {
@@ -202,6 +200,9 @@ func loadLinksIntoMemory() error {
 	}
 	defer rows.Close()
 
+	linkMapMutex.Lock()
+	defer linkMapMutex.Unlock()
+
 	for rows.Next() {
 		var id, url string
 		var clicks int
diff --git a/docker-compose.yml b/docker-compose.yml
index a802854..3e91ed4 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -22,9 +22,10 @@ services:
       POSTGRES_DB: shortr
     ports:
       - "5432:5432"
+    restart: unless-stopped
     volumes:
       - pgdata:/var/lib/postgresql/data
 
-
 volumes:
   pgdata:
+    driver: local