summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorsinanmohd <sinan@sinanmohd.com>2024-07-06 13:56:26 +0530
committersinanmohd <sinan@sinanmohd.com>2024-07-06 13:56:26 +0530
commit4c8be274a50104917fd777645b2985b92eecdfb5 (patch)
tree90769a802835645aeca344e1e99b5effa5b5f44b
parent2f0483eb34ae51dc0dc241409342db021b995856 (diff)
bpf: update database before exit
-rw-r--r--bpf/main.go21
1 files changed, 16 insertions, 5 deletions
diff --git a/bpf/main.go b/bpf/main.go
index c436534..b465f2c 100644
--- a/bpf/main.go
+++ b/bpf/main.go
@@ -5,6 +5,9 @@ import (
"errors"
"log"
"net"
+ "os"
+ "os/signal"
+ "syscall"
"time"
"github.com/cilium/ebpf"
@@ -54,6 +57,9 @@ func Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) {
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:
@@ -61,9 +67,14 @@ func Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) {
if err != nil {
log.Printf("updating usageMap: %s", err)
}
-
+ case <-sigs:
+ err := usageMap.updateDb(queries, ctxDb, false)
+ if err != nil {
+ log.Printf("updating Database: %s", err)
+ }
+ os.Exit(0)
case <-dbTicker.C:
- err := usageMap.updateDb(queries, ctxDb)
+ err := usageMap.updateDb(queries, ctxDb, true)
if err != nil {
log.Printf("updating Database: %s", err)
}
@@ -85,11 +96,11 @@ func (usageStat UsageStat) expired(timeStart *time.Time) bool {
return false
}
-func (usageMap UsageMap) updateDb(queries *db.Queries, ctxDb context.Context) error {
+func (usageMap UsageMap) updateDb(queries *db.Queries, ctxDb context.Context, ifExpired bool) error {
timeStart := time.Now()
for key, value := range usageMap {
- if !value.expired(&timeStart) {
+ if ifExpired && !value.expired(&timeStart) {
continue
}
@@ -149,7 +160,7 @@ func (usageMap UsageMap) update(ingress *ebpf.Map, egress *ebpf.Map) error {
if errors.Is(err, ebpf.ErrKeyNotExist) {
break
} else if err != nil {
- return err;
+ return err
}
}