From 4ef5f7b379b37baeb3b6de1999be16d089e29738 Mon Sep 17 00:00:00 2001
From: sinanmohd <sinan@sinanmohd.com>
Date: Sat, 6 Jul 2024 15:53:34 +0530
Subject: usage: don't handle signals directly

---
 cmd/redq/main.go | 17 ++++++++++++++++-
 usage/main.go    | 27 ++++++++-------------------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/cmd/redq/main.go b/cmd/redq/main.go
index 8807105..df275d6 100644
--- a/cmd/redq/main.go
+++ b/cmd/redq/main.go
@@ -4,10 +4,13 @@ import (
 	"context"
 	"log"
 	"net"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/jackc/pgx/v5"
-	"sinanmohd.com/redq/usage"
 	"sinanmohd.com/redq/db"
+	"sinanmohd.com/redq/usage"
 )
 
 func main() {
@@ -15,6 +18,7 @@ func main() {
 		Data : make(usage.UsageMap),
 	}
 
+
 	iface, err := net.InterfaceByName("wlan0")
 	if err != nil {
 		log.Fatalf("lookup network: %s", err)
@@ -28,5 +32,16 @@ func main() {
 	defer conn.Close(ctx)
 	queries := db.New(conn)
 
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, os.Interrupt, os.Kill, syscall.SIGTERM)
+	go func ()  {
+		<- sigs
+		err := u.UpdateDb(queries, ctx, false)
+		if err != nil {
+			log.Printf("updating Database: %s", err)
+		}
+		os.Exit(0)
+	}()
+
 	u.Run(iface, queries, ctx)
 }
diff --git a/usage/main.go b/usage/main.go
index 8015aa0..2293281 100644
--- a/usage/main.go
+++ b/usage/main.go
@@ -5,10 +5,7 @@ import (
 	"errors"
 	"log"
 	"net"
-	"os"
-	"os/signal"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/cilium/ebpf"
@@ -30,18 +27,18 @@ type UsageMap map[uint64]UsageStat
 type Usage struct {
 	Data UsageMap
 	Mutex sync.Mutex
+	objs bpfObjects
 }
 
 func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) {
-	objs := bpfObjects{}
-	if err := loadBpfObjects(&objs, nil); err != nil {
+	if err := loadBpfObjects(&u.objs, nil); err != nil {
 		log.Fatalf("loading objects: %s", err)
 	}
-	defer objs.Close()
+	defer u.objs.Close()
 
 	ingressLink, err := link.AttachTCX(link.TCXOptions{
 		Interface: iface.Index,
-		Program:   objs.IngressFunc,
+		Program:   u.objs.IngressFunc,
 		Attach:    ebpf.AttachTCXIngress,
 	})
 	if err != nil {
@@ -51,7 +48,7 @@ func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Con
 
 	egressLink, err := link.AttachTCX(link.TCXOptions{
 		Interface: iface.Index,
-		Program:   objs.EgressFunc,
+		Program:   u.objs.EgressFunc,
 		Attach:    ebpf.AttachTCXEgress,
 	})
 	if err != nil {
@@ -63,24 +60,16 @@ func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Con
 	defer bpfTicker.Stop()
 	dbTicker := time.NewTicker(time.Minute)
 	defer dbTicker.Stop()
-	sigs := make(chan os.Signal, 1)
-	signal.Notify(sigs, os.Interrupt, os.Kill, syscall.SIGTERM)
 
 	for {
 		select {
 		case <-bpfTicker.C:
-			err := u.update(objs.IngressIp4UsageMap, objs.EgressIp4UsageMap)
+			err := u.update(u.objs.IngressIp4UsageMap, u.objs.EgressIp4UsageMap)
 			if err != nil {
 				log.Printf("updating usageMap: %s", err)
 			}
-		case <-sigs:
-			err := u.updateDb(queries, ctxDb, false)
-			if err != nil {
-				log.Printf("updating Database: %s", err)
-			}
-			os.Exit(0)
 		case <-dbTicker.C:
-			err := u.updateDb(queries, ctxDb, true)
+			err := u.UpdateDb(queries, ctxDb, true)
 			if err != nil {
 				log.Printf("updating Database: %s", err)
 			}
@@ -102,7 +91,7 @@ func (usageStat *UsageStat) expired(timeStart *time.Time) bool {
 	return false
 }
 
-func (u *Usage) updateDb(queries *db.Queries, ctxDb context.Context, ifExpired bool) error {
+func (u *Usage) UpdateDb(queries *db.Queries, ctxDb context.Context, ifExpired bool) error {
 	timeStart := time.Now()
 
 	u.Mutex.Lock()
-- 
cgit v1.2.3