diff options
| -rw-r--r-- | api/main.go | 19 | ||||
| -rw-r--r-- | cmd/redq/main.go | 16 | ||||
| -rw-r--r-- | usage/main.go | 61 |
3 files changed, 49 insertions, 47 deletions
diff --git a/api/main.go b/api/main.go index 5ce05c1..69cf981 100644 --- a/api/main.go +++ b/api/main.go @@ -13,20 +13,21 @@ type Api struct { sock net.Listener } -func (a *Api) Init() error { +func Close(a *Api) { + a.sock.Close() +} + +func New() (*Api, error) { var err error + var a Api a.sock, err = net.Listen("unix", sockPath) if err != nil { log.Printf("listening on unix socket: %s", err) - return err + return nil, err } - return nil -} - -func handleConn(conn net.Conn) { - defer conn.Close() + return &a, nil } func (a *Api) Run(u *usage.Usage) { @@ -41,6 +42,6 @@ func (a *Api) Run(u *usage.Usage) { } } -func (a *Api) CleanUp() { - a.sock.Close() +func handleConn(conn net.Conn) { + defer conn.Close() } diff --git a/cmd/redq/main.go b/cmd/redq/main.go index f368dae..288c6be 100644 --- a/cmd/redq/main.go +++ b/cmd/redq/main.go @@ -9,14 +9,14 @@ import ( "syscall" "github.com/jackc/pgx/v5" - "sinanmohd.com/redq/db" "sinanmohd.com/redq/api" + "sinanmohd.com/redq/db" "sinanmohd.com/redq/usage" ) func main() { - var u usage.Usage - var a api.Api + var u *usage.Usage + var a *api.Api iface, err := net.InterfaceByName("wlan0") if err != nil { @@ -31,11 +31,11 @@ func main() { defer conn.Close(ctx) queries := db.New(conn) - err = a.Init() + a, err = api.New() if err != nil { os.Exit(0) } - err = u.Init(iface) + u, err = usage.New(iface) if err != nil { os.Exit(0) } @@ -44,11 +44,11 @@ func main() { signal.Notify(sigs, os.Interrupt, os.Kill, syscall.SIGTERM) go func() { <-sigs - u.CleanUp(queries, ctx) - a.CleanUp() + usage.Close(u, queries, ctx) + api.Close(a) os.Exit(0) }() go u.Run(iface, queries, ctx) - a.Run(&u) + a.Run(u) } diff --git a/usage/main.go b/usage/main.go index cc16311..72cf4fd 100644 --- a/usage/main.go +++ b/usage/main.go @@ -31,12 +31,24 @@ type Usage struct { egressLink, ingressLink link.Link } -func (u *Usage) Init(iface *net.Interface) error { +func Close(u *Usage, 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 New(iface *net.Interface) (*Usage, error) { var err error + var u Usage if err := loadBpfObjects(&u.objs, nil); err != nil { log.Printf("loading objects: %s", err) - return err + return nil, err } defer func() { if err != nil { @@ -51,7 +63,7 @@ func (u *Usage) Init(iface *net.Interface) error { }) if err != nil { log.Printf("could not attach TCx program: %s", err) - return err + return nil, err } defer func() { if err != nil { @@ -66,22 +78,11 @@ func (u *Usage) Init(iface *net.Interface) error { }) if err != nil { log.Printf("could not attach TCx program: %s", err) - return err + return nil, err } 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() + return &u, nil } func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) { @@ -106,20 +107,6 @@ func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Con } } -func (us *usageStat) expired(timeStart *time.Time) bool { - timeDiff := timeStart.Sub(us.lastSeen) - if timeDiff > time.Minute { - return true - } - - timeDiff = timeStart.Sub(us.lastDbPush) - if timeDiff > time.Hour { - return true - } - - return false -} - func (u *Usage) UpdateDb(queries *db.Queries, ctxDb context.Context, ifExpired bool) error { timeStart := time.Now() @@ -153,6 +140,20 @@ func (u *Usage) UpdateDb(queries *db.Queries, ctxDb context.Context, ifExpired b return nil } +func (us *usageStat) expired(timeStart *time.Time) bool { + timeDiff := timeStart.Sub(us.lastSeen) + if timeDiff > time.Minute { + return true + } + + timeDiff = timeStart.Sub(us.lastDbPush) + if timeDiff > time.Hour { + return true + } + + return false +} + func (u *Usage) update(ingress *ebpf.Map, egress *ebpf.Map) error { timeStart := time.Now() batchKeys := make([]uint64, 4096) |