feat(info): track gpud process self resource usage (file descriptors, RSS, start time, db size)

components/info/component.go

@@ -3,28 +3,38 @@ package info
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"net"
+	"os"
 	"time"
 
 	"github.com/leptonai/gpud/components"
 	info_id "github.com/leptonai/gpud/components/info/id"
+	"github.com/leptonai/gpud/components/state"
 	"github.com/leptonai/gpud/log"
 	"github.com/leptonai/gpud/manager"
+	"github.com/leptonai/gpud/pkg/file"
+	"github.com/leptonai/gpud/pkg/memory"
+	"github.com/leptonai/gpud/pkg/uptime"
 	"github.com/leptonai/gpud/version"
+
+	"github.com/dustin/go-humanize"
 )
 
-func New(annotations map[string]string) components.Component {
+func New(annotations map[string]string, dbRO *sql.DB) components.Component {
 	return &component{
 		annotations: annotations,
+		dbRO:        dbRO,
 	}
 }
 
 var _ components.Component = (*component)(nil)
 
 type component struct {
 	annotations map[string]string
+	dbRO        *sql.DB
 }
 
 func (c *component) Name() string { return info_id.Name }
@@ -36,6 +46,19 @@ const (
 	StateKeyMacAddress    = "mac_address"
 	StateKeyPackages      = "packages"
 
+	StateKeyGPUdPID = "gpud_pid"
+
+	StateKeyGPUdUsageFileDescriptors = "gpud_usage_file_descriptors"
+
+	StateKeyGPUdUsageMemoryInBytes   = "gpud_usage_memory_in_bytes"
+	StateKeyGPUdUsageMemoryHumanized = "gpud_usage_memory_humanized"
+
+	StateKeyGPUdUsageDBInBytes   = "gpud_usage_db_in_bytes"
+	StateKeyGPUdUsageDBHumanized = "gpud_usage_db_humanized"
+
+	StateKeyGPUdStartTimeInUnixTime = "gpud_start_time_in_unix_time"
+	StateKeyGPUdStartTimeHumanized  = "gpud_start_time_humanized"
+
 	StateNameAnnotations = "annotations"
 )
 
@@ -63,6 +86,37 @@ func (c *component) States(ctx context.Context) ([]components.State, error) {
 		rawPayload, _ := json.Marshal(&packageStatus)
 		managedPackages = string(rawPayload)
 	}
+
+	pid := os.Getpid()
+	gpudUsageFileDescriptors, err := file.GetCurrentProcessUsage()
+	if err != nil {
+		return nil, err
+	}
+
+	gpudUsageMemoryInBytes, err := memory.GetCurrentProcessRSSInBytes()
+	if err != nil {
+		return nil, err
+	}
+	gpudUsageMemoryHumanized := humanize.Bytes(gpudUsageMemoryInBytes)
+
+	var (
+		dbSize          uint64
+		dbSizeHumanized string
+	)
+	if c.dbRO != nil {
+		dbSize, err = state.ReadDBSize(ctx, c.dbRO)
+		if err != nil {
+			return nil, err
+		}
+		dbSizeHumanized = humanize.Bytes(dbSize)
+	}
+
+	gpudStartTimeInUnixTime, err := uptime.GetCurrentProcessStartTimeInUnixTime()
+	if err != nil {
+		return nil, err
+	}
+	gpudStartTimeHumanized := humanize.Time(time.Unix(int64(gpudStartTimeInUnixTime), 0))
+
 	return []components.State{
 		{
 			Name:    StateNameDaemon,
@@ -72,6 +126,19 @@ func (c *component) States(ctx context.Context) ([]components.State, error) {
 				StateKeyDaemonVersion: version.Version,
 				StateKeyMacAddress:    mac,
 				StateKeyPackages:      managedPackages,
+
+				StateKeyGPUdPID: fmt.Sprintf("%d", pid),
+
+				StateKeyGPUdUsageFileDescriptors: fmt.Sprintf("%d", gpudUsageFileDescriptors),
+
+				StateKeyGPUdUsageMemoryInBytes:   fmt.Sprintf("%d", gpudUsageMemoryInBytes),
+				StateKeyGPUdUsageMemoryHumanized: gpudUsageMemoryHumanized,
+
+				StateKeyGPUdUsageDBInBytes:   fmt.Sprintf("%d", dbSize),
+				StateKeyGPUdUsageDBHumanized: dbSizeHumanized,
+
+				StateKeyGPUdStartTimeInUnixTime: fmt.Sprintf("%d", gpudStartTimeInUnixTime),
+				StateKeyGPUdStartTimeHumanized:  gpudStartTimeHumanized,
 			},
 		},
 		{

components/info/component_test.go

@@ -8,7 +8,7 @@ import (
 func TestComponent(t *testing.T) {
 	t.Parallel()
 
-	component := New(map[string]string{"a": "b"})
+	component := New(map[string]string{"a": "b"}, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

components/state/state.go

@@ -166,14 +166,6 @@ UPDATE %s SET %s = '%s' WHERE %s = '%s';
 }
 
 var (
-	currentPages = prometheus.NewGauge(
-		prometheus.GaugeOpts{
-			Namespace: "gpud",
-			Subsystem: "state_sqlite",
-			Name:      "current_pages",
-			Help:      "current number of pages",
-		},
-	)
 	currentSize = prometheus.NewGauge(
 		prometheus.GaugeOpts{
 			Namespace: "gpud",
@@ -185,36 +177,40 @@ var (
 )
 
 func Register(reg *prometheus.Registry) error {
-	if err := reg.Register(currentPages); err != nil {
-		return err
-	}
 	if err := reg.Register(currentSize); err != nil {
 		return err
 	}
 	return nil
 }
 
-// Requires read-write db instance.
-func RecordMetrics(ctx context.Context, db *sql.DB) error {
+func ReadDBSize(ctx context.Context, db *sql.DB) (uint64, error) {
 	var pageCount uint64
 	err := db.QueryRowContext(ctx, "PRAGMA page_count").Scan(&pageCount)
 	if err == sql.ErrNoRows {
-		return errors.New("no page count")
+		return 0, errors.New("no page count")
 	}
 	if err != nil {
-		return err
+		return 0, err
 	}
-	currentPages.Set(float64(pageCount))
 
 	var pageSize uint64
 	err = db.QueryRowContext(ctx, "PRAGMA page_size").Scan(&pageSize)
 	if err == sql.ErrNoRows {
-		return errors.New("no page size")
+		return 0, errors.New("no page size")
+	}
+	if err != nil {
+		return 0, err
 	}
+
+	return pageCount * pageSize, nil
+}
+
+func RecordMetrics(ctx context.Context, db *sql.DB) error {
+	dbSize, err := ReadDBSize(ctx, db)
 	if err != nil {
 		return err
 	}
-	currentSize.Set(float64(pageCount * pageSize))
+	currentSize.Set(float64(dbSize))
 
 	return nil
 }

components/state/state_test.go

@@ -6,6 +6,8 @@ import (
 	"time"
 
 	"github.com/leptonai/gpud/pkg/sqlite"
+
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 func TestOpenMemory(t *testing.T) {
@@ -46,16 +48,36 @@ func TestOpenMemory(t *testing.T) {
 func TestRecordMetrics(t *testing.T) {
 	t.Parallel()
 
-	dbRW, dbRO, close := sqlite.OpenTestDB(t)
-	defer close()
+	reg := prometheus.NewRegistry()
+	if err := Register(reg); err != nil {
+		t.Fatal("failed to register metrics:", err)
+	}
+
+	dbRW, dbRO, cleanup := sqlite.OpenTestDB(t)
+	defer cleanup()
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	if err := RecordMetrics(ctx, dbRO); err == nil {
-		t.Fatal("expected error but got nil")
+	if err := CreateTableMachineMetadata(ctx, dbRW); err != nil {
+		t.Fatal("failed to create table:", err)
+	}
+	id, err := CreateMachineIDIfNotExist(ctx, dbRW, dbRW, "")
+	if err != nil {
+		t.Fatal("failed to create machine id:", err)
 	}
-	if err := RecordMetrics(ctx, dbRW); err != nil {
+	t.Log(id)
+
+	if err := RecordMetrics(ctx, dbRO); err != nil {
 		t.Fatal("failed to record metrics:", err)
 	}
+	if err := Compact(ctx, dbRW); err != nil {
+		t.Fatal("failed to compact database:", err)
+	}
+
+	size, err := ReadDBSize(ctx, dbRO)
+	if err != nil {
+		t.Fatal("failed to read db size:", err)
+	}
+	t.Log(size)
 }

internal/server/server.go

@@ -649,7 +649,7 @@ func New(ctx context.Context, config *lepconfig.Config, endpoint string, cliUID
 			}
 
 		case info_id.Name:
-			allComponents = append(allComponents, info.New(config.Annotations))
+			allComponents = append(allComponents, info.New(config.Annotations, dbRO))
 
 		case memory_id.Name:
 			cfg := memory.Config{Query: defaultQueryCfg}

pkg/file/descriptors_darwin.go

@@ -28,3 +28,9 @@ func GetFileHandles() (uint64, uint64, error) {
 func GetUsage() (uint64, error) {
 	return 0, nil
 }
+
+// Returns the number of allocated file handles for the current process.
+// Not implemented for this architecture.
+func GetCurrentProcessUsage() (uint64, error) {
+	return 0, nil
+}

pkg/file/descriptors_linux.go

@@ -79,3 +79,16 @@ func GetUsage() (uint64, error) {
 	}
 	return total, nil
 }
+
+// Returns the number of allocated file handles for the current process.
+func GetCurrentProcessUsage() (uint64, error) {
+	proc, err := procfs.Self()
+	if err != nil {
+		return 0, err
+	}
+	fdLen, err := proc.FileDescriptorsLen()
+	if err != nil {
+		return 0, err
+	}
+	return uint64(fdLen), nil
+}

pkg/file/descriptors_test.go

@@ -37,3 +37,11 @@ func Test_getFileHandles(t *testing.T) {
 		t.Fatalf("unused is not 0: %v", unused)
 	}
 }
+
+func TestGetCurrentProcessUsage(t *testing.T) {
+	usage, err := GetCurrentProcessUsage()
+	if err != nil {
+		t.Fatalf("failed to get current process usage: %v", err)
+	}
+	t.Logf("usage: %v", usage)
+}

pkg/file/descriptors_windows.go

@@ -22,3 +22,9 @@ func GetFileHandles() (uint64, uint64, error) {
 func GetUsage() (uint64, error) {
 	return 0, nil
 }
+
+// Returns the number of allocated file handles for the current process.
+// Not implemented for this architecture.
+func GetCurrentProcessUsage() (uint64, error) {
+	return 0, nil
+}

pkg/memory/memory.go

@@ -0,0 +1,2 @@
+// Package memory provides utilities for memory usage.
+package memory

pkg/memory/memory_darwin.go

@@ -0,0 +1,8 @@
+//go:build darwin
+// +build darwin
+
+package memory
+
+func GetCurrentProcessRSSInBytes() (uint64, error) {
+	return 0, nil
+}

pkg/memory/memory_linux.go

@@ -0,0 +1,34 @@
+// Copyright 2019 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build linux
+// +build linux
+
+package memory
+
+import "github.com/prometheus/procfs"
+
+// Returns the number of allocated memory RSS usage in bytes for the current process.
+func GetCurrentProcessRSSInBytes() (uint64, error) {
+	proc, err := procfs.Self()
+	if err != nil {
+		return 0, err
+	}
+
+	stat, err := proc.Stat()
+	if err != nil {
+		return 0, err
+	}
+
+	return uint64(stat.ResidentMemory()), nil
+}

pkg/memory/memory_test.go

@@ -0,0 +1,13 @@
+package memory
+
+import (
+	"testing"
+)
+
+func TestGetCurrentProcessRSSInBytes(t *testing.T) {
+	bytes, err := GetCurrentProcessRSSInBytes()
+	if err != nil {
+		t.Fatalf("failed to get bytes: %v", err)
+	}
+	t.Logf("bytes: %v", bytes)
+}

pkg/memory/memory_windows.go

@@ -0,0 +1,8 @@
+//go:build windows
+// +build windows
+
+package memory
+
+func GetCurrentProcessRSSInBytes() (uint64, error) {
+	return 0, nil
+}

pkg/uptime/uptime.go

@@ -0,0 +1,2 @@
+// Package uptime provides utilities for uptime.
+package uptime

pkg/uptime/uptime_darwin.go

@@ -0,0 +1,8 @@
+//go:build darwin
+// +build darwin
+
+package uptime
+
+func GetCurrentProcessStartTimeInUnixTime() (uint64, error) {
+	return 0, nil
+}