summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorsinanmohd <sinan@sinanmohd.com>2024-07-06 15:53:34 +0530
committersinanmohd <sinan@sinanmohd.com>2024-07-06 15:55:08 +0530
commit4ef5f7b379b37baeb3b6de1999be16d089e29738 (patch)
tree48a8fcf17ff70e645544a7a9ca27e6ab0b4eeb8b
parent9e4f86953d7a39e4d32801a328c1f8424ca6c969 (diff)
usage: don't handle signals directly
-rw-r--r--cmd/redq/main.go17
-rw-r--r--usage/main.go27
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()