Add NetworkLimiter with bidirectional bandwidth limiting via tc integration

README.md

@@ -1,7 +1,7 @@
 # Giogo
 
 Giogo is a command-line tool that allows you to run processes with specified resource limitations using Linux cgroups.  
-It provides an easy-to-use interface to limit CPU, memory, and IO resources for a process and its children.
+It provides an easy-to-use interface to limit CPU, memory, IO, and network resources for a process and its children.
 
 **Note: Root privileges are required, and cgroups v1 is currently not supported.**
 
@@ -16,13 +16,15 @@ It provides an easy-to-use interface to limit CPU, memory, and IO resources for
   - [CPU Limitations](#cpu-limitations)
   - [Memory Limitations](#memory-limitations)
   - [IO Limitations](#io-limitations)
+  - [Network Limitations](#network-limitations)
 - [Examples](#examples)
 
 ## Features
 
 - **CPU Limiting**: Restrict CPU usage as a fraction of total CPU time.
 - **Memory Limiting**: Set maximum memory usage.
 - **IO Limiting**: Control IO read and write bandwidth.
+- **Network Limiting**: Set network class identifiers and priorities for network traffic.
 - **Cgroups Support**: Works with cgroups v2 only (cgroups v1 is not supported at this time).
 - **Process Isolation**: Limits apply to the process and all its child processes.
 
@@ -122,6 +124,41 @@ By default, Giogo sets a bandwidth throttle on every block device's IO. The Linu
 **Additional Note:**  
 If your operations utilize the `O_DIRECT` flag, the RAM limit is not required, as `O_DIRECT` bypasses the kernel's caching mechanism.
 
+### Network Limitations
+
+- **`--network-class-id=VALUE`**
+
+  Set a class identifier for the container's network packets.
+
+  - **`VALUE`**: A numeric identifier (uint32) used for packet classification and traffic control.
+  - **Example**: `--network-class-id=100` sets the network class identifier to 100.
+  - **Use Case**: This can be used in conjunction with Linux traffic control (tc) for advanced network QoS configuration.
+
+- **`--network-priority=VALUE`**
+
+  Set the priority of network traffic for the container.
+
+  - **`VALUE`**: A numeric priority value (uint32), where higher values typically indicate higher priority.
+  - **Example**: `--network-priority=50` sets the network traffic priority to 50.
+  - **Use Case**: Helps prioritize network traffic when multiple processes compete for bandwidth.
+
+- **`--network-max-bandwidth=VALUE`**
+
+  Set a maximum network bandwidth limit for the container. Requires `--network-class-id` to be set.
+
+  - **`VALUE`**: Maximum bandwidth using the same notation as memory (`k`, `m`, `g`).
+  - **Units**:
+    - `k` or `K`: Kilobytes per second
+    - `m` or `M`: Megabytes per second
+    - `g` or `G`: Gigabytes per second
+  - **Example**: `--network-max-bandwidth=1m` limits network bandwidth to 1 MB/s.
+  - **Use Case**: Enforces hard bandwidth limits on network traffic using Linux traffic control (tc) with HTB qdisc.
+
+**Note:**  
+Network limitations work with cgroups v2's network controller to provide packet classification and prioritization. The priority setting applies to all network interfaces in the container.
+
+When `--network-max-bandwidth` is specified with `--network-class-id`, giogo automatically configures Linux traffic control (tc) with HTB (Hierarchical Token Bucket) to enforce the bandwidth limit. The tc rules are automatically cleaned up when the process exits.
+
 ## Examples
 
 ### Limit CPU and Memory
@@ -147,3 +184,27 @@ sudo giogo --io-read-max=2m --io-write-max=1m --ram=2g -- your_io_intensive_comm
 ```
 
 - **Description**: Runs `your_io_intensive_command` with IO read limited to 2 MB/s and IO write limited to 1 MB/s, while setting a high RAM limit of 2 GB to bypass the default association between `io-write-max` and RAM usage.
+
+### Network Traffic Control
+
+```bash
+sudo giogo --network-class-id=100 --network-priority=50 -- your_network_intensive_app
+```
+
+- **Description**: Runs `your_network_intensive_app` with network class identifier set to 100 and network priority set to 50, allowing for packet classification and traffic prioritization.
+
+### Network Bandwidth Limiting
+
+```bash
+sudo giogo --network-class-id=100 --network-max-bandwidth=1m -- your_app
+```
+
+- **Description**: Runs `your_app` with network bandwidth limited to 1 MB/s. This automatically configures traffic control (tc) with HTB qdisc to enforce the limit.
+
+### Combined Resource Limitation
+
+```bash
+sudo giogo --cpu=0.5 --ram=512m --network-class-id=200 --network-max-bandwidth=500k -- your_app
+```
+
+- **Description**: Runs `your_app` with CPU limited to 50% of one core, RAM limited to 512 MB, network class identifier set to 200, and network bandwidth limited to 500 KB/s.

internal/cli/cli.go

@@ -11,10 +11,13 @@ import (
 )
 
 var (
-	ram        string
-	cpu        string
-	ioReadMax  string
-	ioWriteMax string
+	ram                 string
+	cpu                 string
+	ioReadMax           string
+	ioWriteMax          string
+	networkClassID      string
+	networkPriority     string
+	networkMaxBandwidth string
 )
 
 func SetupRootCommand(rootCmd *cobra.Command) {
@@ -28,6 +31,9 @@ func SetupRootCommand(rootCmd *cobra.Command) {
 	rootCmd.Flags().StringVar(&cpu, "cpu", "", "CPU limit as a fraction between 0 and 1 (e.g., 0.5)")
 	rootCmd.Flags().StringVar(&ioReadMax, "io-read-max", limiter.UnlimitedIOValue, "IO read max bandwidth (e.g., 128k, 1m)")
 	rootCmd.Flags().StringVar(&ioWriteMax, "io-write-max", limiter.UnlimitedIOValue, "IO write max bandwidth (e.g., 128k, 1m)")
+	rootCmd.Flags().StringVar(&networkClassID, "network-class-id", "", "Network class identifier for container's network packets")
+	rootCmd.Flags().StringVar(&networkPriority, "network-priority", "", "Network priority for container's network traffic")
+	rootCmd.Flags().StringVar(&networkMaxBandwidth, "network-max-bandwidth", "", "Maximum network bandwidth (e.g., 1m, 500k) - requires network-class-id")
 }
 
 func Execute() {
@@ -41,7 +47,7 @@ func Execute() {
 }
 
 // TODO: This logic should be moved to a separate package as it's part of the core functionality
-func CreateLimiters(cpu, ram, ioReadMax, ioWriteMax string) ([]limiter.ResourceLimiter, error) {
+func CreateLimiters(cpu, ram, ioReadMax, ioWriteMax, networkClassID, networkPriority, networkMaxBandwidth string) ([]limiter.ResourceLimiter, error) {
 	var limiters []limiter.ResourceLimiter
 
 	if cpu != "" {
@@ -89,11 +95,24 @@ func CreateLimiters(cpu, ram, ioReadMax, ioWriteMax string) ([]limiter.ResourceL
 		// Known issue: a minimum amount of memory is required to start a process, so if the memory limit is too low, the process will not start.
 	}
 
+	if networkClassID != "" || networkPriority != "" || networkMaxBandwidth != "" {
+		netInit := limiter.NetworkLimiterInitializer{
+			ClassID:      networkClassID,
+			Priority:     networkPriority,
+			MaxBandwidth: networkMaxBandwidth,
+		}
+		netLimiter, err := limiter.NewNetworkLimiter(&netInit)
+		if err != nil {
+			return nil, fmt.Errorf("invalid network value: %v", err)
+		}
+		limiters = append(limiters, netLimiter)
+	}
+
 	return limiters, nil
 }
 
 func runCommand(cmd *cobra.Command, args []string) error {
-	limiters, err := CreateLimiters(cpu, ram, ioReadMax, ioWriteMax)
+	limiters, err := CreateLimiters(cpu, ram, ioReadMax, ioWriteMax, networkClassID, networkPriority, networkMaxBandwidth)
 	if err != nil {
 		return err
 	}

internal/cli/cli_test.go

@@ -68,7 +68,7 @@ func TestCreateLimiters_IOLimitsAndNoRAM(t *testing.T) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax)
+	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax, "", "", "")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -99,7 +99,7 @@ func TestCreateLimiters_IOLimitsAndNoRAM_ReadLimit(t *testing.T) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax)
+	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax, "", "", "")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -130,7 +130,7 @@ func TestCreateLimiters_IOLimitsAndRAM(t *testing.T) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax)
+	limiters, err := cli.CreateLimiters(cpu, ram, ioReadMax, ioWriteMax, "", "", "")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}

internal/executor/executor.go

@@ -1,6 +1,9 @@
 package executor
 
 import (
+	"fmt"
+	"os"
+
 	"github.com/pmarchini/giogo/internal/core"
 	"github.com/pmarchini/giogo/internal/limiter"
 
@@ -23,6 +26,35 @@ func (e *Executor) RunCommand(args []string) error {
 		l.Apply(&resources)
 	}
 
+	// Track successfully setup lifecycle limiters for cleanup
+	var setupLimiters []limiter.LifecycleLimiter
+
+	// Setup any limiters that implement lifecycle methods
+	for _, l := range e.Limiters {
+		if lifecycle, ok := l.(limiter.LifecycleLimiter); ok {
+			if err := lifecycle.Setup(); err != nil {
+				// Cleanup already setup limiters before returning error
+				for i := len(setupLimiters) - 1; i >= 0; i-- {
+					if cleanupErr := setupLimiters[i].Cleanup(); cleanupErr != nil {
+						fmt.Fprintf(os.Stderr, "Warning: failed to cleanup limiter during error recovery: %v\n", cleanupErr)
+					}
+				}
+				return fmt.Errorf("failed to setup limiter: %v", err)
+			}
+			setupLimiters = append(setupLimiters, lifecycle)
+		}
+	}
+
+	// Ensure cleanup happens for all successfully setup limiters
+	defer func() {
+		for i := len(setupLimiters) - 1; i >= 0; i-- {
+			if err := setupLimiters[i].Cleanup(); err != nil {
+				// Non-fatal: log the error but don't fail the command
+				fmt.Fprintf(os.Stderr, "Warning: failed to cleanup limiter: %v\n", err)
+			}
+		}
+	}()
+
 	coreModule, err := core.NewCore(resources)
 	if err != nil {
 		return err

internal/limiter/limiter.go

@@ -7,3 +7,11 @@ import (
 type ResourceLimiter interface {
 	Apply(resources *specs.LinuxResources)
 }
+
+// LifecycleLimiter is an optional interface for limiters that need
+// setup and cleanup operations beyond the cgroup configuration
+type LifecycleLimiter interface {
+	ResourceLimiter
+	Setup() error
+	Cleanup() error
+}