diff options
-rw-r--r-- | api/main.go | 10 | ||||
-rw-r--r-- | api/usage.go | 48 | ||||
-rw-r--r-- | cmd/redq/main.go | 2 | ||||
-rw-r--r-- | db/query.sql | 3 | ||||
-rw-r--r-- | db/query.sql.go | 16 |
5 files changed, 75 insertions, 4 deletions
diff --git a/api/main.go b/api/main.go index 7081d3c..32f7d08 100644 --- a/api/main.go +++ b/api/main.go @@ -1,10 +1,12 @@ package api import ( + "context" "encoding/json" "log" "net" + "sinanmohd.com/redq/db" "sinanmohd.com/redq/usage" ) @@ -40,7 +42,7 @@ func New() (*Api, error) { return &a, nil } -func (a *Api) Run(u *usage.Usage) { +func (a *Api) Run(u *usage.Usage, queries *db.Queries, ctxDb context.Context) { for { conn, err := a.sock.Accept() if err != nil { @@ -48,11 +50,11 @@ func (a *Api) Run(u *usage.Usage) { continue } - go handleConn(conn, u) + go handleConn(conn, u, queries, ctxDb) } } -func handleConn(conn net.Conn, u *usage.Usage) { +func handleConn(conn net.Conn, u *usage.Usage, queries *db.Queries, ctxDb context.Context) { defer conn.Close() var req ApiReq buf := make([]byte, bufSize) @@ -72,6 +74,8 @@ func handleConn(conn net.Conn, u *usage.Usage) { switch req.Type { case "bandwidth": handleBandwidth(conn, u) + case "usage": + handleUsage(conn, u, queries, ctxDb) default: log.Printf("invalid request type: %s", req.Type) } diff --git a/api/usage.go b/api/usage.go new file mode 100644 index 0000000..5849ba7 --- /dev/null +++ b/api/usage.go @@ -0,0 +1,48 @@ +package api + +import ( + "context" + "encoding/json" + "log" + "net" + + "github.com/dustin/go-humanize" + "sinanmohd.com/redq/db" + "sinanmohd.com/redq/usage" +) + +type UsageStat struct { + Ingress string `json:"ingress"` + Egress string `json:"egress"` +} + +type UsageResp map[string]UsageStat + +func handleUsage(conn net.Conn, u *usage.Usage, queries *db.Queries, ctxDb context.Context) { + resp := make(UsageResp) + + fetchedUsage, err := queries.GetUsage(ctxDb) + if err != nil { + log.Printf("fetching from database: %s", err) + return + } + + u.Mutex.RLock() + for _, value := range u.Data { + fetchedUsage.Ingress += int64(value.Ingress) + fetchedUsage.Egress += int64(value.Egress) + } + u.Mutex.RUnlock() + resp["total"] = UsageStat{ + Ingress: humanize.Bytes(uint64(fetchedUsage.Ingress)), + Egress: humanize.Bytes(uint64(fetchedUsage.Egress)), + } + + buf, err := json.Marshal(resp) + if err != nil { + log.Printf("marshaling json: %s", err) + return + } + + conn.Write(buf) +} diff --git a/cmd/redq/main.go b/cmd/redq/main.go index 288c6be..f9d9542 100644 --- a/cmd/redq/main.go +++ b/cmd/redq/main.go @@ -50,5 +50,5 @@ func main() { }() go u.Run(iface, queries, ctx) - a.Run(u) + a.Run(u, queries, ctx) } diff --git a/db/query.sql b/db/query.sql index 75d5b61..cfea3f1 100644 --- a/db/query.sql +++ b/db/query.sql @@ -4,3 +4,6 @@ INSERT INTO Usage ( ) VALUES ( $1, $2, $3, $4, $5 ); + +-- name: GetUsage :one +SELECT SUM(Ingress) AS Ingress, SUM(Egress) AS Egress FROM Usage; diff --git a/db/query.sql.go b/db/query.sql.go index 7afe159..d6304a7 100644 --- a/db/query.sql.go +++ b/db/query.sql.go @@ -37,3 +37,19 @@ func (q *Queries) EnterUsage(ctx context.Context, arg EnterUsageParams) error { ) return err } + +const getUsage = `-- name: GetUsage :one +SELECT SUM(Ingress) AS Ingress, SUM(Egress) AS Egress FROM Usage +` + +type GetUsageRow struct { + Ingress int64 + Egress int64 +} + +func (q *Queries) GetUsage(ctx context.Context) (GetUsageRow, error) { + row := q.db.QueryRow(ctx, getUsage) + var i GetUsageRow + err := row.Scan(&i.Ingress, &i.Egress) + return i, err +} |