bpf: update database before exit

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
 		}
 	}