actually migrate code for new structure

pkg/roles/dhcp/api_leases.go

@@ -47,11 +47,16 @@ func (r *Role) APILeasesGet() usecase.Interactor {
 			r.log.Warn("failed to get scope", zap.Error(err))
 			return status.Wrap(errors.New("failed to get scope"), status.Internal)
 		}
+		s, err := r.scopeFromKV(rawScope.Kvs[0])
+		if err != nil {
+			r.log.Warn("failed to parse scope", zap.Error(err))
+			return status.Wrap(err, status.Internal)
+		}
 
 		leaseKey := r.i.KV().Key(
 			types.KeyRole,
 			types.KeyScopes,
-			input.ScopeName,
+			s.Name,
 		)
 		if input.Identifier == "" {
 			leaseKey = leaseKey.Prefix(true)
@@ -63,7 +68,7 @@ func (r *Role) APILeasesGet() usecase.Interactor {
 			return status.Wrap(err, status.Internal)
 		}
 		for _, lease := range rawLeases.Kvs {
-			l, err := r.leaseFromKV(lease)
+			l, err := s.leaseFromKV(lease)
 			if err != nil {
 				r.log.Warn("failed to parse lease", zap.Error(err))
 				continue
@@ -127,7 +132,7 @@ func (r *Role) APILeasesPut() usecase.Interactor {
 			return status.Wrap(errors.New("failed to construct scope"), status.Internal)
 		}
 
-		l := r.NewLease(input.Identifier)
+		l := scope.NewLease(input.Identifier)
 		l.Address = input.Address
 		l.Hostname = input.Hostname
 		l.AddressLeaseTime = input.AddressLeaseTime
@@ -156,13 +161,40 @@ type APILeasesWOLInput struct {
 
 func (r *Role) APILeasesWOL() usecase.Interactor {
 	u := usecase.NewInteractor(func(ctx context.Context, input APILeasesWOLInput, output *struct{}) error {
-		r.leasesM.RLock()
-		l, ok := r.leases[input.Identifier]
-		r.leasesM.RUnlock()
-		if !ok {
-			return status.InvalidArgument
+		rawScope, err := r.i.KV().Get(
+			ctx,
+			r.i.KV().Key(
+				types.KeyRole,
+				types.KeyScopes,
+				input.Scope,
+			).String(),
+		)
+		if err != nil || len(rawScope.Kvs) < 1 {
+			r.log.Warn("failed to get scope", zap.Error(err))
+			return status.Wrap(errors.New("failed to get scope"), status.Internal)
 		}
-		err := l.sendWOL()
+		scope, err := r.scopeFromKV(rawScope.Kvs[0])
+		if err != nil {
+			r.log.Warn("failed to construct scope", zap.Error(err))
+			return status.Wrap(errors.New("failed to construct scope"), status.Internal)
+		}
+
+		leaseKey := r.i.KV().Key(
+			types.KeyRole,
+			types.KeyScopes,
+			scope.Name,
+			input.Identifier,
+		)
+		rawLeases, err := r.i.KV().Get(ctx, leaseKey.String(), clientv3.WithPrefix())
+		if err != nil || len(rawLeases.Kvs) < 1 {
+			return status.Wrap(err, status.InvalidArgument)
+		}
+		l, err := scope.leaseFromKV(rawLeases.Kvs[0])
+		if err != nil {
+			return status.Wrap(err, status.Internal)
+		}
+
+		err = l.sendWOL()
 		if err != nil {
 			return status.Wrap(err, status.Internal)
 		}

pkg/roles/dhcp/api_leases_test.go

@@ -33,7 +33,7 @@ func TestAPILeasesGet(t *testing.T) {
 			types.KeyScopes,
 			scope.Name,
 		).String(),
-		tests.MustJSON(scope),
+		tests.MustJSON(&scope),
 	))
 	lease := testLease()
 	tests.PanicIfError(inst.KV().Put(
@@ -70,7 +70,7 @@ func TestAPILeasesPut(t *testing.T) {
 			types.KeyScopes,
 			scope.Name,
 		).String(),
-		tests.MustJSON(scope),
+		tests.MustJSON(&scope),
 	))
 	assert.NoError(t, role.APILeasesPut().Interact(ctx, dhcp.APILeasesPutInput{
 		Identifier: name,
@@ -111,7 +111,7 @@ func TestAPILeasesDelete(t *testing.T) {
 			types.KeyScopes,
 			scope.Name,
 		).String(),
-		tests.MustJSON(scope),
+		tests.MustJSON(&scope),
 	))
 	lease := testLease()
 	tests.PanicIfError(inst.KV().Put(

pkg/roles/dhcp/ipam_internal.go

@@ -147,13 +147,9 @@ func (i *InternalIPAM) IsIPFree(ip netip.Addr, identifier *string) bool {
 		return false
 	}
 	// check for existing leases
-	i.role.leasesM.RLock()
-	defer i.role.leasesM.RUnlock()
-	for _, l := range i.role.leases {
-		// Ignore leases from other scopes
-		if l.ScopeKey != i.scope.Name {
-			continue
-		}
+	i.scope.leasesSync.RLock()
+	defer i.scope.leasesSync.RUnlock()
+	for _, l := range i.scope.leases {
 		if l.Address != ip.String() {
 			continue
 		}

pkg/roles/dhcp/ipam_internal_test.go

@@ -58,7 +58,7 @@ func TestIPAMInternal_NextFreeAddress_UniqueParallel(t *testing.T) {
 			types.KeyScopes,
 			scope.Name,
 		).String(),
-		tests.MustJSON(scope),
+		tests.MustJSON(&scope),
 	))
 	// Create fake leases to test against
 	for i := 0; i < iter-10; i++ {

pkg/roles/dhcp/leases.go

@@ -5,7 +5,6 @@ import (
 	"encoding/base64"
 	"encoding/hex"
 	"encoding/json"
-	"fmt"
 	"math"
 	"net"
 	"net/netip"
@@ -42,47 +41,24 @@ type Lease struct {
 }
 
 func (r *Role) FindLease(req *Request4) *Lease {
-	r.leasesM.RLock()
-	defer r.leasesM.RUnlock()
-	lease, ok := r.leases[r.DeviceIdentifier(req.DHCPv4)]
+	expectedScope := r.findScopeForRequest(req)
+	expectedScope.leasesSync.RLock()
+	defer expectedScope.leasesSync.RUnlock()
+	lease, ok := expectedScope.leases[r.DeviceIdentifier(req.DHCPv4)]
 	if !ok {
 		return nil
 	}
-	// Check if the leases's scope matches the expected scope to handle this request
-	expectedScope := r.findScopeForRequest(req, func(scope *Scope) int {
-		// Consider the existing lease for finding the scope
-		// Check how many bits of the leases address match the scope
-		sm := scope.match(net.ParseIP(lease.Address))
-		// If the matching bits match how many bits are in the CIDR, and
-		// the scope of the lease matches the scope we're filtering, we've
-		// got a match
-		if sm == lease.scope.cidr.Bits() && lease.ScopeKey == scope.Name {
-			return 99
-		}
-		return -1
-	})
-	if expectedScope != nil && lease.scope != expectedScope {
-		// We have a specific scope to handle this request but it doesn't match the lease
-		lease.scope = expectedScope
-		lease.ScopeKey = expectedScope.Name
-		lease.setLeaseIP(req)
-		lease.log.Info("Re-assigning address for lease due to changed request scope", zap.String("newIP", lease.Address))
-		go func() {
-			err := lease.Put(req.Context, lease.scope.TTL)
-			if err != nil {
-				r.log.Warn("failed to update lease", zap.Error(err))
-			}
-		}()
-	}
 	return lease
 }
 
-func (r *Role) NewLease(identifier string) *Lease {
+func (s *Scope) NewLease(identifier string) *Lease {
 	return &Lease{
-		inst:       r.i,
+		inst:       s.inst,
 		Identifier: identifier,
-		log:        r.log.With(zap.String("identifier", identifier)),
+		log:        s.log.With(zap.String("identifier", identifier)),
 		Expiry:     0,
+		scope:      s,
+		ScopeKey:   s.Name,
 	}
 }
 
@@ -107,27 +83,21 @@ func (l *Lease) setLeaseIP(req *Request4) {
 	l.scope.ipam.UseIP(*ip, l.Identifier)
 }
 
-func (r *Role) leaseFromKV(raw *mvccpb.KeyValue) (*Lease, error) {
-	prefix := r.i.KV().Key(
+func (s *Scope) leaseFromKV(raw *mvccpb.KeyValue) (*Lease, error) {
+	prefix := s.inst.KV().Key(
 		types.KeyRole,
 		types.KeyScopes,
 	).Prefix(true).String()
 	keyParts := strings.SplitN(prefix, "/", 2)
 	identifier := strings.TrimPrefix(string(raw.Key), prefix+"/"+keyParts[0])
-	l := r.NewLease(identifier)
+	l := s.NewLease(identifier)
 	err := json.Unmarshal(raw.Value, &l)
 	if err != nil {
 		return l, err
 	}
 	l.etcdKey = string(raw.Key)
-
-	r.scopesM.RLock()
-	scope, ok := r.scopes[l.ScopeKey]
-	r.scopesM.RUnlock()
-	if !ok {
-		return l, fmt.Errorf("DHCP lease with invalid scope key: %s", l.ScopeKey)
-	}
-	l.scope = scope
+	l.scope = s
+	l.ScopeKey = s.Name
 	return l, nil
 }
 

pkg/roles/dhcp/metrics.go

@@ -1,8 +1,6 @@
 package dhcp
 
 import (
-	"math/big"
-
 	"beryju.io/gravity/pkg/extconfig"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -36,14 +34,8 @@ var (
 func (s *Scope) calculateUsage() {
 	usable := s.ipam.UsableSize()
 	dhcpScopeSize.WithLabelValues(s.Name).Set(float64(usable.Uint64()))
-	used := big.NewInt(0)
-	s.role.leasesM.RLock()
-	defer s.role.leasesM.RUnlock()
-	for _, lease := range s.role.leases {
-		if lease.ScopeKey != s.Name {
-			continue
-		}
-		used = used.Add(used, big.NewInt(1))
-	}
-	dhcpScopeUsage.WithLabelValues(s.Name).Set(float64(used.Uint64()))
+	s.leasesSync.RLock()
+	defer s.leasesSync.RUnlock()
+	used := len(s.leases)
+	dhcpScopeUsage.WithLabelValues(s.Name).Set(float64(used))
 }

pkg/roles/dhcp/role.go

@@ -26,7 +26,6 @@ type Role struct {
 	ctx context.Context
 
 	scopes map[string]*Scope
-	leases map[string]*Lease
 
 	cfg *RoleConfig
 
@@ -35,7 +34,6 @@ type Role struct {
 
 	oui     *oui.OuiDb
 	scopesM sync.RWMutex
-	leasesM sync.RWMutex
 }
 
 func New(instance roles.Instance) *Role {
@@ -44,8 +42,6 @@ func New(instance roles.Instance) *Role {
 		i:       instance,
 		scopes:  make(map[string]*Scope),
 		scopesM: sync.RWMutex{},
-		leases:  make(map[string]*Lease),
-		leasesM: sync.RWMutex{},
 		ctx:     instance.Context(),
 	}
 	r.s4 = &handler4{
@@ -78,16 +74,8 @@ func (r *Role) Start(ctx context.Context, config []byte) error {
 	start := sentry.TransactionFromContext(ctx).StartChild("gravity.dhcp.start")
 	defer start.Finish()
 	r.loadInitialScopes(start.Context())
-	r.loadInitialLeases(start.Context())
-
-	// Since scope usage relies on r.leases, but r.leases is loaded after the scopes,
-	// manually update the usage
-	for _, s := range r.scopes {
-		s.calculateUsage()
-	}
 
 	go r.startWatchScopes()
-	go r.startWatchLeases()
 
 	err := r.initServer4()
 	if err != nil {