Improve the StorageLocal struct

packages/api/internal/handlers/template_build_status.go

@@ -109,6 +109,7 @@ func (a *APIStore) GetTemplatesTemplateIDBuildsBuildIDStatus(c *gin.Context, tem
 	if err != nil {
 		telemetry.ReportError(ctx, "error when comparing versions", err, telemetry.WithTemplateID(templateID), telemetry.WithBuildID(buildID))
 		a.sendAPIStoreError(c, http.StatusInternalServerError, "Error when processing build logs")
+
 		return
 	}
 

packages/api/internal/sandbox/store/memory/reservation.go

@@ -46,16 +46,19 @@ func (s *ReservationStorage) Reserve(teamID, sandboxID string, limit int64) (fin
 		if sbx, ok := teamSandboxes[sandboxID]; ok {
 			alreadyPresent = true
 			startResult = sbx.start
+
 			return teamSandboxes
 		}
 
 		if limit >= 0 && len(teamSandboxes) >= int(limit) {
 			limitExceeded = true
+
 			return teamSandboxes
 		}
 
 		startResult = utils.NewSetOnce[sandbox.Sandbox]()
 		teamSandboxes[sandboxID] = newSandboxReservation(startResult)
+
 		return teamSandboxes
 	})
 
@@ -87,6 +90,7 @@ func (s *ReservationStorage) Remove(teamID, sandboxID string) {
 		}
 
 		delete(ts, sandboxID)
+
 		return len(ts) == 0
 	})
 }

packages/api/internal/sandbox/store/memory/reservation_test.go

@@ -411,6 +411,7 @@ func TestReservation_ConcurrentWaitAndFinish(t *testing.T) {
 			}
 
 			waiters[i] = waitForStart
+
 			return nil
 		})
 	}

packages/orchestrator/internal/cfg/model.go

@@ -1,6 +1,8 @@
 package cfg
 
 import (
+	"net"
+	"reflect"
 	"time"
 
 	"github.com/caarlos0/env/v11"
@@ -35,9 +37,17 @@ type Config struct {
 }
 
 func Parse() (Config, error) {
-	return env.ParseAs[Config]()
+	return env.ParseAsWithOptions[Config](env.Options{
+		FuncMap: map[reflect.Type]env.ParserFunc{
+			reflect.TypeOf(net.IPNet{}): network.ParseIPNet,
+		},
+	})
 }
 
 func ParseBuilder() (BuilderConfig, error) {
-	return env.ParseAs[BuilderConfig]()
+	return env.ParseAsWithOptions[BuilderConfig](env.Options{
+		FuncMap: map[reflect.Type]env.ParserFunc{
+			reflect.TypeOf(net.IPNet{}): network.ParseIPNet,
+		},
+	})
 }

packages/orchestrator/internal/sandbox/network/config.go

@@ -0,0 +1,52 @@
+package network
+
+import (
+	"fmt"
+	"net"
+	"reflect"
+
+	"github.com/caarlos0/env/v11"
+)
+
+type Config struct {
+	// Using reserver IPv4 in range that is used for experiments and documentation
+	// https://en.wikipedia.org/wiki/Reserved_IP_addresses
+	HyperloopIPAddress       string `env:"SANDBOX_HYPERLOOP_IP"         envDefault:"192.0.2.1"`
+	HyperloopProxyPort       uint16 `env:"SANDBOX_HYPERLOOP_PROXY_PORT" envDefault:"5010"`
+	UseLocalNamespaceStorage bool   `env:"USE_LOCAL_NAMESPACE_STORAGE"`
+	LocalNamespaceStorageDir string `env:"LOCAL_NAMESPACE_STORAGE_DIR"  envDefault:"/run/orchestrator/netns"`
+
+	SandboxesHostNetworkCIDR    *net.IPNet `env:"SANDBOXES_HOST_NETWORK_CIDR"    envDefault:"10.11.0.0/16"`
+	SandboxesVirtualNetworkCIDR *net.IPNet `env:"SANDBOXES_VIRTUAL_NETWORK_CIDR" envDefault:"10.12.0.0/16"`
+}
+
+func (c Config) GetVirtualSlotsSize() int {
+	ones, _ := c.SandboxesVirtualNetworkCIDR.Mask.Size()
+
+	// total IPs in the CIDR block
+	totalIPs := 1 << (32 - ones)
+
+	// total slots that we can allocate
+	// we need to divide total IPs by number of addresses per slot (vpeer and veth)
+	// then we subtract the number of addresses so it will not overflow, because we are adding them incrementally by slot index
+	totalSlots := (totalIPs / vrtAddressPerSlot) - vrtAddressPerSlot
+
+	return totalSlots
+}
+
+func ParseConfig() (Config, error) {
+	return env.ParseAsWithOptions[Config](env.Options{
+		FuncMap: map[reflect.Type]env.ParserFunc{
+			reflect.TypeOf(net.IPNet{}): ParseIPNet,
+		},
+	})
+}
+
+func ParseIPNet(cidr string) (any, error) {
+	_, subnet, err := net.ParseCIDR(cidr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse network CIDR %s: %w", cidr, err)
+	}
+
+	return *subnet, nil
+}

packages/orchestrator/internal/sandbox/network/config_test.go

@@ -0,0 +1,32 @@
+package network
+
+import (
+	"net"
+	"reflect"
+	"testing"
+
+	"github.com/caarlos0/env/v11"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseIPNet(t *testing.T) {
+	t.Setenv("CIDR_OVERRIDE_DEFAULT", "5.0.0.0/8")
+
+	var testConfig struct {
+		HasDefault           *net.IPNet `env:"CIDR_USE_DEFAULT"             envDefault:"10.10.0.0/16"`
+		OverrideDefault      *net.IPNet `env:"CIDR_OVERRIDE_DEFAULT"        envDefault:"1.2.3.4/16"`
+		NilBecauseNonePassed *net.IPNet `env:"CIDR_NIL_BECAUSE_NONE_PASSED"`
+	}
+
+	err := env.ParseWithOptions(&testConfig, env.Options{
+		FuncMap: map[reflect.Type]env.ParserFunc{
+			reflect.TypeOf(net.IPNet{}): ParseIPNet,
+		},
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, "10.10.0.0/16", testConfig.HasDefault.String())
+	assert.Equal(t, "5.0.0.0/8", testConfig.OverrideDefault.String())
+	assert.Nil(t, testConfig.NilBecauseNonePassed)
+}

packages/orchestrator/internal/sandbox/network/pool.go

@@ -6,8 +6,6 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/caarlos0/env/v11"
-	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 	"go.uber.org/zap"
@@ -22,8 +20,6 @@ const (
 )
 
 var (
-	meter = otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network")
-
 	newSlotsAvailableCounter = utils.Must(meter.Int64UpDownCounter("orchestrator.network.slots_pool.new",
 		metric.WithDescription("Number of new network slots ready to be used."),
 		metric.WithUnit("{slot"),
@@ -46,18 +42,6 @@ var (
 	))
 )
 
-type Config struct {
-	// Using reserver IPv4 in range that is used for experiments and documentation
-	// https://en.wikipedia.org/wiki/Reserved_IP_addresses
-	HyperloopIPAddress       string `env:"SANDBOX_HYPERLOOP_IP"         envDefault:"192.0.2.1"`
-	HyperloopProxyPort       uint16 `env:"SANDBOX_HYPERLOOP_PROXY_PORT" envDefault:"5010"`
-	UseLocalNamespaceStorage bool   `env:"USE_LOCAL_NAMESPACE_STORAGE"`
-}
-
-func ParseConfig() (Config, error) {
-	return env.ParseAs[Config]()
-}
-
 type Pool struct {
 	config Config
 
@@ -76,7 +60,7 @@ func NewPool(newSlotsPoolSize, reusedSlotsPoolSize int, nodeID string, config Co
 	newSlots := make(chan *Slot, newSlotsPoolSize-1)
 	reusedSlots := make(chan *Slot, reusedSlotsPoolSize)
 
-	slotStorage, err := NewStorage(vrtSlotsSize, nodeID, config)
+	slotStorage, err := NewStorage(nodeID, config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create slot storage: %w", err)
 	}
@@ -100,7 +84,7 @@ func (p *Pool) createNetworkSlot(ctx context.Context) (*Slot, error) {
 
 	err = ips.CreateNetwork()
 	if err != nil {
-		releaseErr := p.slotStorage.Release(ips)
+		releaseErr := p.slotStorage.Release(ctx, ips)
 		err = errors.Join(err, releaseErr)
 
 		return nil, fmt.Errorf("failed to create network: %w", err)
@@ -212,7 +196,7 @@ func (p *Pool) cleanup(ctx context.Context, slot *Slot) error {
 		errs = append(errs, fmt.Errorf("cannot remove network when releasing slot '%d': %w", slot.Idx, err))
 	}
 
-	err = p.slotStorage.Release(slot)
+	err = p.slotStorage.Release(ctx, slot)
 	if err != nil {
 		errs = append(errs, fmt.Errorf("failed to release slot '%d': %w", slot.Idx, err))
 	}

packages/orchestrator/internal/sandbox/network/slot.go

@@ -3,7 +3,6 @@ package network
 import (
 	"context"
 	"fmt"
-	"log"
 	"net"
 	"path/filepath"
 	"strconv"
@@ -14,16 +13,13 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	netutils "k8s.io/utils/net"
-
-	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 )
 
+const netNamespacesDir = "/var/run/netns"
+
 var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network")
 
 const (
-	defaultHostNetworkCIDR = "10.11.0.0/16"
-	defaultVrtNetworkCIDR  = "10.12.0.0/16"
-
 	hostMask          = 32
 	vrtMask           = 31                  // 2 usable ips per block (vpeer and veth)
 	vrtAddressPerSlot = 1 << (32 - vrtMask) // vrt addresses per slot (vpeer and veth)
@@ -34,12 +30,6 @@ const (
 	tapMAC           = "02:FC:00:00:00:05"
 )
 
-var (
-	hostNetworkCIDR = getHostNetworkCIDR()
-	vrtNetworkCIDR  = getVrtNetworkCIDR()
-	vrtSlotsSize    = GetVrtSlotsSize()
-)
-
 // Slot network allocation
 //
 // For each slot, we allocate three IP addresses:
@@ -55,8 +45,8 @@ var (
 // Vpeer receives the first IP in the block, and Veth receives the second IP. Block is calculated as (slot index * addresses per slot allocation).
 // Vrt address per slot is always 2, so we can allocate /31 CIDR block for each slot.
 type Slot struct {
-	Key string
-	Idx int
+	Name string
+	Idx  int
 
 	Firewall *Firewall
 
@@ -79,11 +69,13 @@ type Slot struct {
 	hyperloopIP, hyperloopPort string
 }
 
-func NewSlot(key string, idx int, config Config) (*Slot, error) {
-	if idx < 1 || idx > vrtSlotsSize {
+func NewSlot(name string, idx int, config Config) (*Slot, error) {
+	if vrtSlotsSize := config.GetVirtualSlotsSize(); idx < 1 || idx > vrtSlotsSize {
 		return nil, fmt.Errorf("slot index %d is out of range [1, %d)", idx, vrtSlotsSize)
 	}
 
+	vrtNetworkCIDR := config.SandboxesVirtualNetworkCIDR
+
 	vEthIp, err := netutils.GetIndexedIP(vrtNetworkCIDR, idx*vrtAddressPerSlot)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get veth indexed IP: %w", err)
@@ -100,7 +92,7 @@ func NewSlot(key string, idx int, config Config) (*Slot, error) {
 		return nil, fmt.Errorf("failed to parse vrt CIDR: %w", err)
 	}
 
-	hostIp, err := netutils.GetIndexedIP(hostNetworkCIDR, idx)
+	hostIp, err := netutils.GetIndexedIP(config.SandboxesHostNetworkCIDR, idx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get host IP: %w", err)
 	}
@@ -118,8 +110,8 @@ func NewSlot(key string, idx int, config Config) (*Slot, error) {
 	}
 
 	slot := &Slot{
-		Key: key,
-		Idx: idx,
+		Name: name,
+		Idx:  idx,
 
 		vPeerIp: vPeerIp,
 		vEthIp:  vEthIp,
@@ -223,7 +215,7 @@ func (s *Slot) TapMAC() string {
 
 func (s *Slot) InitializeFirewall() error {
 	if s.Firewall != nil {
-		return fmt.Errorf("firewall is already initialized for slot %s", s.Key)
+		return fmt.Errorf("firewall is already initialized for slot %s", s.Name)
 	}
 
 	fw, err := NewFirewall(s.TapName())
@@ -313,45 +305,3 @@ func (s *Slot) ResetInternet(ctx context.Context) error {
 
 	return nil
 }
-
-func getHostNetworkCIDR() *net.IPNet {
-	cidr := env.GetEnv("SANDBOXES_HOST_NETWORK_CIDR", defaultHostNetworkCIDR)
-
-	_, subnet, err := net.ParseCIDR(cidr)
-	if err != nil {
-		log.Fatalf("Failed to parse network CIDR %s: %v", cidr, err)
-	}
-
-	log.Println("Using host network cidr", "cidr", cidr)
-
-	return subnet
-}
-
-func getVrtNetworkCIDR() *net.IPNet {
-	cidr := env.GetEnv("SANDBOXES_VRT_NETWORK_CIDR", defaultVrtNetworkCIDR)
-
-	_, subnet, err := net.ParseCIDR(cidr)
-	if err != nil {
-		log.Fatalf("Failed to parse network CIDR %s: %v", cidr, err)
-	}
-
-	log.Printf("Using vrt network cidr %s", cidr)
-
-	return subnet
-}
-
-func GetVrtSlotsSize() int {
-	ones, _ := getVrtNetworkCIDR().Mask.Size()
-
-	// total IPs in the CIDR block
-	totalIPs := 1 << (32 - ones)
-
-	// total slots that we can allocate
-	// we need to divide total IPs by number of addresses per slot (vpeer and veth)
-	// then we subtract the number of addresses so it will not overflow, because we are adding them incrementally by slot index
-	totalSlots := (totalIPs / vrtAddressPerSlot) - vrtAddressPerSlot
-
-	log.Printf("Using network slot size: %d", totalSlots)
-
-	return totalSlots
-}

packages/orchestrator/internal/sandbox/network/storage.go

@@ -8,14 +8,14 @@ import (
 
 type Storage interface {
 	Acquire(ctx context.Context) (*Slot, error)
-	Release(s *Slot) error
+	Release(ctx context.Context, s *Slot) error
 }
 
 // NewStorage creates a new slot storage based on the environment, we are ok with using a memory storage for local
-func NewStorage(slotsSize int, nodeID string, config Config) (Storage, error) {
+func NewStorage(nodeID string, config Config) (Storage, error) {
 	if env.IsDevelopment() || config.UseLocalNamespaceStorage {
-		return NewStorageLocal(slotsSize, config)
+		return NewStorageLocal(config)
 	}
 
-	return NewStorageKV(slotsSize, nodeID, config)
+	return NewStorageKV(nodeID, config)
 }