Project policy

core2/pkg/accounting/00_module.go

@@ -43,6 +43,9 @@ func Init() {
 	initGrantsExport()
 	times["GrantsExport"] = t.Mark()
 
+	initPolicySubscriptions()
+	times["PolicySubscriptions"] = t.Mark()
+
 	coreutil.PrintStartupTimes("Accounting", times)
 
 	if util.DevelopmentModeEnabled() {

core2/pkg/accounting/policy_cache.go

@@ -0,0 +1,76 @@
+package accounting
+
+import (
+	"context"
+	"sync"
+
+	"ucloud.dk/core/pkg/coreutil"
+	db "ucloud.dk/shared/pkg/database"
+	fndapi "ucloud.dk/shared/pkg/foundation"
+)
+
+// policyCache is a mapping of projectId -> map[schemaName] -> PolicySpecification
+var policyCache struct {
+	Mu                sync.RWMutex
+	PoliciesByProject map[string]map[fndapi.PolicyName]fndapi.Specification
+}
+
+func initPolicySubscriptions() {
+
+	policyCache.Mu.Lock()
+	policyCache.PoliciesByProject = make(map[string]map[fndapi.PolicyName]fndapi.Specification)
+	policyCache.Mu.Unlock()
+
+	go func() {
+		policyUpdates := db.Listen(context.Background(), "policy_updates")
+		policyDeletes := db.Listen(context.Background(), "policy_deleted")
+
+		var projectId string
+		var policySpecifications map[fndapi.PolicyName]fndapi.Specification
+		var policiesOk bool
+
+		for {
+			select {
+			case projectId = <-policyUpdates:
+				db.NewTx0(func(tx *db.Transaction) {
+					policySpecifications, policiesOk = coreutil.PolicySpecificationsRetrieveFromDatabase(tx, projectId)
+				})
+			case projectId = <-policyDeletes:
+				db.NewTx0(func(tx *db.Transaction) {
+
+					policySpecifications, policiesOk = coreutil.PolicySpecificationsRetrieveFromDatabase(tx, projectId)
+				})
+			}
+
+			if policiesOk {
+				updatePolicyCacheForProject(projectId, policySpecifications)
+			}
+		}
+
+	}()
+}
+
+// policiesByProject returns mapping of [schema Name] => PolicySpecification. If no policy is cached for the project it
+// will attempt to retrieve it from DB. This is also how it is populated.
+func policiesByProject(projectId string) map[fndapi.PolicyName]fndapi.Specification {
+	policyCache.Mu.Lock()
+	projectPolicies, ok := policyCache.PoliciesByProject[projectId]
+	if !ok {
+		db.NewTx0(func(tx *db.Transaction) {
+			policySpecifications, policiesOk := coreutil.PolicySpecificationsRetrieveFromDatabase(tx, projectId)
+			if policiesOk {
+				policyCache.PoliciesByProject[projectId] = policySpecifications
+			}
+			projectPolicies = policySpecifications
+		})
+	}
+	policyCache.Mu.Unlock()
+
+	return projectPolicies
+}
+
+func updatePolicyCacheForProject(projectId string, policySpecifications map[fndapi.PolicyName]fndapi.Specification) {
+	policyCache.Mu.Lock()
+	policyCache.PoliciesByProject[projectId] = policySpecifications
+	policyCache.Mu.Unlock()
+}

core2/pkg/accounting/provider_notifications.go

@@ -23,20 +23,68 @@ import (
 var providerWalletNotifications = make(chan AccWalletId, 1024*1024)
 
 var providerNotifications struct {
-	Mu                        sync.Mutex
+	Mu sync.Mutex
+
+	// All of these follow providerId -> sessionId -> channel
+
 	ProjectChannelsByProvider map[string]map[string]chan *fndapi.Project
 	WalletsByProvider         map[string]map[string]chan *accapi.WalletV2
+	PoliciesByProvider        map[string]map[string]chan fndapi.PoliciesForProject
+}
+
+func retrieveRelevantProviders(projectId string) map[string]util.Empty {
+	projectWallets := internalRetrieveWallets(time.Now(), projectId, walletFilter{RequireActive: true})
+	relevantProviders := map[string]util.Empty{}
+
+	for _, w := range projectWallets {
+		relevantProviders[w.PaysFor.Provider] = util.Empty{}
+	}
+
+	return relevantProviders
 }
 
 func initProviderNotifications() {
 	providerNotifications.ProjectChannelsByProvider = map[string]map[string]chan *fndapi.Project{}
 	providerNotifications.WalletsByProvider = map[string]map[string]chan *accapi.WalletV2{}
+	providerNotifications.PoliciesByProvider = map[string]map[string]chan fndapi.PoliciesForProject{}
+
+	policyCache.Mu.Lock()
+	policyCache.PoliciesByProject = make(map[string]map[fndapi.PolicyName]fndapi.Specification)
+	policyCache.Mu.Unlock()
 
 	go func() {
 		// NOTE(Dan): These two channels receive events from database triggers set on the relevant insert/update/delete
 		// operations. The payload is either the project or group ID which triggered the update.
 		projectUpdates := db.Listen(context.Background(), "project_updates")
 		groupUpdates := db.Listen(context.Background(), "project_group_updates")
+		policyUpdates := db.Listen(context.Background(), "policy_updates")
+
+		broadcastPoliciesForProject := func(projectId string, policySpecifications map[fndapi.PolicyName]fndapi.Specification) {
+			relevantProviders := retrieveRelevantProviders(projectId)
+			var allChannels []chan fndapi.PoliciesForProject
+
+			providerNotifications.Mu.Lock()
+
+			for provider := range relevantProviders {
+				channels, ok := providerNotifications.PoliciesByProvider[provider]
+				if ok {
+					for _, ch := range channels {
+						allChannels = append(allChannels, ch)
+					}
+				}
+			}
+
+			providerNotifications.Mu.Unlock()
+
+			updatePolicyCacheForProject(projectId, policySpecifications)
+
+			for _, ch := range allChannels {
+				select {
+				case ch <- fndapi.PoliciesForProject{ProjectId: projectId, PoliciesByName: policySpecifications}:
+				case <-time.After(200 * time.Millisecond):
+				}
+			}
+		}
 
 		for {
 			var project fndapi.Project
@@ -45,6 +93,10 @@ func initProviderNotifications() {
 			var walletId AccWalletId
 			var walletOk bool
 
+			var policySpecifications map[fndapi.PolicyName]fndapi.Specification
+			var projectIdForPolicies string
+			var policiesOk bool
+
 			select {
 			case projectId := <-projectUpdates:
 				db.NewTx0(func(tx *db.Transaction) {
@@ -58,15 +110,17 @@ func initProviderNotifications() {
 
 			case walletId = <-providerWalletNotifications:
 				walletOk = true
+
+			case projectId := <-policyUpdates:
+				db.NewTx0(func(tx *db.Transaction) {
+					policySpecifications, policiesOk = coreutil.PolicySpecificationsRetrieveFromDatabase(tx, projectId)
+				})
+				projectIdForPolicies = projectId
+
 			}
 
 			if projectOk {
-				projectWallets := internalRetrieveWallets(time.Now(), project.Id, walletFilter{RequireActive: true})
-				relevantProviders := map[string]util.Empty{}
-
-				for _, w := range projectWallets {
-					relevantProviders[w.PaysFor.Provider] = util.Empty{}
-				}
+				relevantProviders := retrieveRelevantProviders(project.Id)
 
 				var allChannels []chan *fndapi.Project
 
@@ -86,8 +140,17 @@ func initProviderNotifications() {
 					case ch <- &project:
 					case <-time.After(200 * time.Millisecond):
 					}
+				}
 
+				// Project membership changed, so newly relevant providers need
+				// the current policy state even if the policies themselves did not change.
+				db.NewTx0(func(tx *db.Transaction) {
+					policySpecifications, policiesOk = coreutil.PolicySpecificationsRetrieveFromDatabase(tx, project.Id)
+				})
+				if policiesOk {
+					broadcastPoliciesForProject(project.Id, policySpecifications)
 				}
+
 			} else if walletOk {
 				wallet, ok := internalRetrieveWallet(time.Now(), walletId, false)
 				if ok && !wallet.PaysFor.FreeToUse && len(wallet.AllocationGroups) != 0 {
@@ -109,6 +172,8 @@ func initProviderNotifications() {
 						}
 					}
 				}
+			} else if policiesOk {
+				broadcastPoliciesForProject(projectIdForPolicies, policySpecifications)
 			}
 		}
 	}()
@@ -175,6 +240,12 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 			RefToCategory map[int]accapi.ProductCategory
 		}
 
+		policies struct {
+			Counter        int
+			ProjectIdToRef map[string]int
+			RefToProjectId map[int]string
+		}
+
 		ctx    context.Context
 		cancel context.CancelFunc
 	)
@@ -188,13 +259,17 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 	productCategories.IdToRef = map[accapi.ProductCategoryIdV2]int{}
 	productCategories.RefToCategory = map[int]accapi.ProductCategory{}
 
+	policies.ProjectIdToRef = map[string]int{}
+	policies.RefToProjectId = map[int]string{}
+
 	ctx, cancel = context.WithCancel(context.Background())
 
 	// Subscription
 	// -----------------------------------------------------------------------------------------------------------------
 	sessionId := util.RandomTokenNoTs(32)
 	projectUpdates := make(chan *fndapi.Project, 128)
 	walletUpdates := make(chan *accapi.WalletV2, 128)
+	policyUpdates := make(chan fndapi.PoliciesForProject, 128)
 
 	{
 		providerNotifications.Mu.Lock()
@@ -213,6 +288,13 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 		}
 		pmap[sessionId] = projectUpdates
 
+		polmap, ok := providerNotifications.PoliciesByProvider[providerId]
+		if !ok {
+			polmap = map[string]chan fndapi.PoliciesForProject{}
+			providerNotifications.PoliciesByProvider[providerId] = polmap
+		}
+		polmap[sessionId] = policyUpdates
+
 		providerNotifications.Mu.Unlock()
 	}
 
@@ -222,6 +304,7 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 		providerNotifications.Mu.Lock()
 		delete(providerNotifications.WalletsByProvider[providerId], sessionId)
 		delete(providerNotifications.ProjectChannelsByProvider[providerId], sessionId)
+		delete(providerNotifications.PoliciesByProvider[providerId], sessionId)
 		providerNotifications.Mu.Unlock()
 	}()
 
@@ -264,6 +347,16 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 				}
 			}
 		}
+
+		updatedPolicies := coreutil.PoliciesListUpdatedAfter(replayFrom)
+		for _, p := range updatedPolicies {
+			select {
+			case <-ctx.Done():
+				return
+			case policyUpdates <- p:
+			}
+		}
+
 	}()
 
 	// Request processing
@@ -344,6 +437,7 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 
 	projectsToSend := map[int]util.Empty{}
 	usersToSend := map[int]util.Empty{}
+	policiesToSend := map[int]fndapi.PoliciesForProject{}
 	var walletsToSend []*accapi.WalletV2
 	categoriesToSend := map[int]util.Empty{}
 
@@ -365,6 +459,25 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 		return ref
 	}
 
+	appendPolicies := func(policiesForProject fndapi.PoliciesForProject, forced bool) int {
+		projectID := policiesForProject.ProjectId
+		ref, ok := policies.ProjectIdToRef[projectID]
+		if !ok {
+			ref, ok = policies.Counter, true
+			policies.ProjectIdToRef[projectID] = ref
+			policies.RefToProjectId[ref] = projectID
+			policiesToSend[ref] = policiesForProject
+
+			policies.Counter++
+		} else if forced {
+			policies.ProjectIdToRef[projectID] = ref
+			policies.RefToProjectId[ref] = projectID
+			policiesToSend[ref] = policiesForProject
+		}
+
+		return ref
+	}
+
 	appendProjectById := func(projectId string) int {
 		ref, ok := projects.ProjectIdToRef[projectId]
 		if !ok {
@@ -439,6 +552,22 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 			out.WriteString(string(projectJson))
 		}
 
+		for policyRef, _ := range policiesToSend {
+			project := policies.RefToProjectId[policyRef]
+
+			policyCache.Mu.Lock()
+			currentPolices := policyCache.PoliciesByProject[project]
+			projectPolicies := fndapi.PoliciesForProject{
+				project,
+				currentPolices,
+			}
+			currentProjectPoliciesJson, _ := json.Marshal(projectPolicies)
+			policyCache.Mu.Unlock()
+			out.WriteU8(opPolicyChange)
+			out.WriteU32(uint32(policyRef))
+			out.WriteString(string(currentProjectPoliciesJson))
+		}
+
 		for categoryRef, _ := range categoriesToSend {
 			category := productCategories.RefToCategory[categoryRef]
 			categoryJson, _ := json.Marshal(category)
@@ -485,7 +614,8 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 			projectsToSend = map[int]util.Empty{}
 			categoriesToSend = map[int]util.Empty{}
 			usersToSend = map[int]util.Empty{}
-			walletsToSend = nil
+			walletsToSend = []*accapi.WalletV2{}
+			policiesToSend = map[int]fndapi.PoliciesForProject{}
 
 			if err != nil {
 				cancel()
@@ -517,17 +647,23 @@ func providerNotificationHandleClient(conn *ws.Conn) {
 					}
 				}
 			}
+
+		case projectPolicies, ok := <-policyUpdates:
+			if ok {
+				appendPolicies(projectPolicies, true)
+			}
 		}
 
 		flush()
 	}
 }
 
 const (
-	opAuth       = 0
-	opWallet     = 1
-	opProject    = 2
-	opCategory   = 3
-	opUser       = 4
-	opReplayUser = 5
+	opAuth         = 0
+	opWallet       = 1
+	opProject      = 2
+	opCategory     = 3
+	opUser         = 4
+	opReplayUser   = 5
+	opPolicyChange = 6
 )