From f2af6ba15466f7ac39f2c12ea29fab5b0c5415c4 Mon Sep 17 00:00:00 2001
From: sinanmohd <sinan@sinanmohd.com>
Date: Sat, 6 Jul 2024 16:13:08 +0530
Subject: usage/Init: init

---
 cmd/redq/main.go | 14 ++++++--------
 usage/main.go    | 47 ++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/cmd/redq/main.go b/cmd/redq/main.go
index df275d6..abac395 100644
--- a/cmd/redq/main.go
+++ b/cmd/redq/main.go
@@ -14,10 +14,7 @@ import (
 )
 
 func main() {
-	u := &usage.Usage {
-		Data : make(usage.UsageMap),
-	}
-
+	var u usage.Usage
 
 	iface, err := net.InterfaceByName("wlan0")
 	if err != nil {
@@ -32,14 +29,15 @@ func main() {
 	defer conn.Close(ctx)
 	queries := db.New(conn)
 
+	err = u.Init(iface)
+	if err != nil {
+		os.Exit(0)
+	}
 	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)
-		}
+		u.CleanUp(queries, ctx)
 		os.Exit(0)
 	}()
 
diff --git a/usage/main.go b/usage/main.go
index 2293281..8be8ed3 100644
--- a/usage/main.go
+++ b/usage/main.go
@@ -28,34 +28,63 @@ type Usage struct {
 	Data UsageMap
 	Mutex sync.Mutex
 	objs bpfObjects
+	egressLink, ingressLink link.Link
 }
 
-func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) {
+func (u *Usage) Init(iface *net.Interface) error {
+	var err error
+
 	if err := loadBpfObjects(&u.objs, nil); err != nil {
-		log.Fatalf("loading objects: %s", err)
+		log.Printf("loading objects: %s", err)
+		return err;
 	}
-	defer u.objs.Close()
+	defer func() {
+		if err != nil {
+			u.objs.Close()
+		}
+	}()
 
-	ingressLink, err := link.AttachTCX(link.TCXOptions{
+	u.ingressLink, err = link.AttachTCX(link.TCXOptions{
 		Interface: iface.Index,
 		Program:   u.objs.IngressFunc,
 		Attach:    ebpf.AttachTCXIngress,
 	})
 	if err != nil {
-		log.Fatalf("could not attach TCx program: %s", err)
+		log.Printf("could not attach TCx program: %s", err)
+		return err
 	}
-	defer ingressLink.Close()
+	defer func() {
+		if err != nil {
+			u.ingressLink.Close()
+		}
+	}()
 
-	egressLink, err := link.AttachTCX(link.TCXOptions{
+	u.egressLink, err = link.AttachTCX(link.TCXOptions{
 		Interface: iface.Index,
 		Program:   u.objs.EgressFunc,
 		Attach:    ebpf.AttachTCXEgress,
 	})
 	if err != nil {
-		log.Fatalf("could not attach TCx program: %s", err)
+		log.Printf("could not attach TCx program: %s", err)
+		return err
 	}
-	defer egressLink.Close()
 
+	u.Data = make(UsageMap)
+	return nil
+}
+
+func (u *Usage) CleanUp(queries *db.Queries, ctxDb context.Context) {
+	err := u.UpdateDb(queries, ctxDb, false)
+	if err != nil {
+		log.Printf("updating Database: %s", err)
+	}
+
+	u.objs.Close()
+	u.ingressLink.Close()
+	u.egressLink.Close()
+}
+
+func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) {
 	bpfTicker := time.NewTicker(time.Second)
 	defer bpfTicker.Stop()
 	dbTicker := time.NewTicker(time.Minute)
-- 
cgit v1.2.3