Merge pull request #44 from uselagoon/calculate-databases

Improve database storage calculation

Author: shreddedbacon

README.md

@@ -24,15 +24,15 @@ Storage-calculator connects to rabbitmq in lagoon-core and publishes to the acti
 		"claims": [{
 			"environment": 1,
 			"persistentStorageClaim": "nginx",
-			"bytesUsed": 1200
+			"kibUsed": 1200
 		}, {
 			"environment": 1,
 			"persistentStorageClaim": "solr",
-			"bytesUsed": 2200
+			"kibUsed": 2200
 		}, {
 			"environment": 1,
 			"persistentStorageClaim": "mariadb",
-			"bytesUsed": 3200
+			"kibUsed": 3200
 		}]
 	}
 }

cmd/main.go

@@ -35,6 +35,8 @@ import (
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 
 	// "github.com/cheshir/go-mq"
+	mariadbv1 "github.com/amazeeio/dbaas-operator/apis/mariadb/v1"
+	postgresv1 "github.com/amazeeio/dbaas-operator/apis/postgres/v1"
 	mq "github.com/cheshir/go-mq/v2"
 	"k8s.io/apimachinery/pkg/runtime"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
@@ -62,6 +64,8 @@ var (
 
 func init() {
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
+	utilruntime.Must(mariadbv1.AddToScheme(scheme))
+	utilruntime.Must(postgresv1.AddToScheme(scheme))
 
 	//+kubebuilder:scaffold:scheme
 }
@@ -122,7 +126,7 @@ func main() {
 		"The retry interval for rabbitmq.")
 	flag.StringVar(&lagoonAppID, "lagoon-app-id", "storage-calculator",
 		"The appID to use that will be sent with messages.")
-	flag.StringVar(&storageCalculatorImage, "storage-calculator-image", "imagecache.amazeeio.cloud/amazeeio/alpine-mysql-client",
+	flag.StringVar(&storageCalculatorImage, "storage-calculator-image", "uselagoon/database-tools",
 		"The image to use for storage-calculator pods.")
 	flag.BoolVar(&exportPrometheusMetrics, "prometheus-metrics", false, "Export prometheus metrics.")
 

config/rbac/role.yaml

@@ -48,3 +48,12 @@ rules:
   verbs:
   - get
   - list
+- apiGroups:
+  - mariadb.amazee.io
+  - postgres.amazee.io
+  resources:
+  - mariadbconsumers
+  - postgresqlconsumers
+  verbs:
+  - list
+  - watch

go.mod

@@ -3,6 +3,7 @@ module github.com/uselagoon/storage-calculator
 go 1.24.0
 
 require (
+	github.com/amazeeio/dbaas-operator v0.4.0
 	github.com/cheshir/go-mq/v2 v2.0.1
 	github.com/go-logr/logr v1.4.3
 	github.com/onsi/ginkgo/v2 v2.25.2

go.sum

@@ -0,0 +1,127 @@
+package storage
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	mariadbv1 "github.com/amazeeio/dbaas-operator/apis/mariadb/v1"
+	postgresv1 "github.com/amazeeio/dbaas-operator/apis/postgres/v1"
+	"github.com/go-logr/logr"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/meta"
+	client "sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+// +kubebuilder:rbac:groups=mariadb.amazee.io;postgres.amazee.io,resources=mariadbconsumers;postgresqlconsumers,verbs=list;watch
+
+// checkDatabasesCreatePods calculates storage usage for databases managed by dbaas-operator.
+// Databases running as a pod are handled by the volumes calculator.
+func (c *Calculator) checkDatabasesCreatePods(
+	ctx context.Context,
+	opLog logr.Logger,
+	namespace corev1.Namespace,
+) (int, error) {
+	opLog = opLog.WithName("Databases")
+	ignoreRegex := c.IgnoreRegex
+	if value, ok := namespace.Labels["lagoon.sh/storageCalculatorIgnoreRegex"]; ok {
+		ignoreRegex = value
+	}
+
+	environmentID := 0
+	if value, ok := namespace.Labels["lagoon.sh/environmentId"]; ok {
+		environmentID, _ = strconv.Atoi(strings.TrimSpace(value))
+	}
+
+	nsListOption := (&client.ListOptions{}).ApplyOptions([]client.ListOption{
+		client.InNamespace(namespace.Name),
+	})
+
+	// Check for managed MariaDB databases.
+	mariadbs := []mariadbv1.MariaDBConsumer{}
+	mariadbList := &mariadbv1.MariaDBConsumerList{}
+	if err := c.Client.List(ctx, mariadbList, nsListOption); err != nil {
+		if !meta.IsNoMatchError(err) {
+			opLog.Error(err, fmt.Sprintf("error getting mariadbconsumer for namespace %s", namespace.Name))
+		}
+	} else {
+		for _, mariadb := range mariadbList.Items {
+			if ignoreRegex != "" {
+				match, _ := regexp.MatchString(ignoreRegex, mariadb.Name)
+				if match {
+					opLog.Info(fmt.Sprintf("ignoring mariadbconsumer %s", mariadb.Name))
+					continue
+				}
+			}
+			mariadbs = append(mariadbs, mariadb)
+		}
+	}
+
+	// Check for managed PostgreSQL databases.
+	postgresdbs := []postgresv1.PostgreSQLConsumer{}
+	postgresList := &postgresv1.PostgreSQLConsumerList{}
+	if err := c.Client.List(ctx, postgresList, nsListOption); err != nil {
+		if !meta.IsNoMatchError(err) {
+			opLog.Error(err, fmt.Sprintf("error getting postgresqlconsumer for namespace %s", namespace.Name))
+		}
+	} else {
+		for _, postgres := range postgresList.Items {
+			if ignoreRegex != "" {
+				match, _ := regexp.MatchString(ignoreRegex, postgres.Name)
+				if match {
+					opLog.Info(fmt.Sprintf("ignoring postgresconsumer %s", postgres.Name))
+					continue
+				}
+			}
+			postgresdbs = append(postgresdbs, postgres)
+		}
+	}
+
+	if len(mariadbs) == 0 && len(postgresdbs) == 0 {
+		opLog.Info(fmt.Sprintf("no databases in %s", namespace.Name))
+		return 0, nil
+	}
+
+	// Calculate storage for all databases.
+	services := databasesCalculatorPod{
+		MariaDB:    mariadbs,
+		PostgreSQL: postgresdbs,
+	}
+	storageClaims, err := c.createDatabasePod(ctx, opLog, namespace, services, environmentID)
+	if err != nil {
+		return 0, fmt.Errorf("error calculating databases storage: %w", err)
+	}
+
+	if len(storageClaims) == 0 {
+		return 0, nil
+	}
+
+	// Report storage sizes for all databases.
+	actionData := ActionEvent{
+		Type:      "updateEnvironmentStorage",
+		EventType: "environmentStorage",
+		Data: ActionData{
+			Claims: storageClaims,
+		},
+		Meta: MetaData{
+			Namespace:   namespace.Name,
+			Project:     namespace.Labels["lagoon.sh/project"],
+			Environment: namespace.Labels["lagoon.sh/environment"],
+		},
+	}
+	opLog.Info(fmt.Sprintf("storage in %s: %v", namespace.Name, actionData))
+
+	if c.ExportMetrics {
+		actionData.ExportMetrics(c.PromStorage)
+	}
+
+	actionDataJSON, _ := json.Marshal(actionData)
+	if err := c.MQ.Publish("lagoon-actions", actionDataJSON); err != nil {
+		return 0, fmt.Errorf("error publishing databases storage to mq: %v", err)
+	}
+
+	return len(storageClaims), nil
+}

internal/storage/databases.go

@@ -2,8 +2,10 @@ package storage
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"strconv"
+	"strings"
 
 	"github.com/go-logr/logr"
 	"github.com/prometheus/client_golang/prometheus"
@@ -71,9 +73,7 @@ type ActionData struct {
 type StorageClaim struct {
 	Environment          int    `json:"environment"`
 	PersisteStorageClaim string `json:"persistentStorageClaim"`
-	// it is actually kilobytes that du outputs, this should be deprecated
-	BytesUsed uint64 `json:"bytesUsed"`
-	KiBUsed   uint64 `json:"kibUsed"`
+	KiBUsed              uint64 `json:"kibUsed"`
 }
 
 // Calculate will run the storage-calculator job.
@@ -100,9 +100,70 @@ func (c *Calculator) Calculate() {
 		return
 	}
 	for _, namespace := range namespaces.Items {
-		err := c.checkVolumesCreatePods(ctx, opLog, namespace)
-		if err != nil {
-			continue
+		opLog.Info(fmt.Sprintf("calculating storage for namespace %s", namespace.Name))
+
+		// Cleanup old calculator pods if they are still running.
+		p1, _ := labels.NewRequirement("lagoon.sh/storageCalculator", "in", []string{"true"})
+		labelRequirements := []labels.Requirement{}
+		labelRequirements = append(labelRequirements, *p1)
+		podListOption := (&client.ListOptions{}).ApplyOptions([]client.ListOption{
+			client.InNamespace(namespace.Name),
+			client.MatchingLabelsSelector{
+				Selector: labels.NewSelector().Add(labelRequirements...),
+			},
+		})
+		pods := &corev1.PodList{}
+		if err := c.Client.List(ctx, pods, podListOption); err != nil {
+			opLog.Error(err, fmt.Sprintf("error getting storage-calculator pods for namespace %s", namespace.Name))
+		}
+		for _, pod := range pods.Items {
+			c.cleanup(ctx, opLog, &pod)
+		}
+
+		// Calculate storage for volumes and databases.
+		volClaims, volErr := c.checkVolumesCreatePods(ctx, opLog, namespace)
+		if volErr != nil {
+			opLog.Error(volErr, "error calculating volumes")
+		}
+		dbClaims, dbErr := c.checkDatabasesCreatePods(ctx, opLog, namespace)
+		if dbErr != nil {
+			opLog.Error(dbErr, "error calculating databases")
+		}
+
+		// Report a "none" claim to differentiate between no storage used and no
+		// storage recorded (due to error or crash)
+		if (volClaims == 0 && volErr != nil) && (dbClaims == 0 && dbErr != nil) {
+			environmentID := 0
+			if value, ok := namespace.Labels["lagoon.sh/environmentId"]; ok {
+				environmentID, _ = strconv.Atoi(strings.TrimSpace(value))
+			}
+
+			storData := ActionData{}
+			storData.Claims = append(storData.Claims, StorageClaim{
+				Environment:          environmentID,
+				PersisteStorageClaim: "none",
+				KiBUsed:              0,
+			})
+			actionData := ActionEvent{
+				Type:      "updateEnvironmentStorage",
+				EventType: "environmentStorage",
+				Data:      storData,
+				Meta: MetaData{
+					Namespace:   namespace.Name,
+					Project:     namespace.Labels["lagoon.sh/project"],
+					Environment: namespace.Labels["lagoon.sh/environment"],
+				},
+			}
+			opLog.Info(fmt.Sprintf("no storage used in %s: %v", namespace.Name, actionData))
+
+			if c.ExportMetrics {
+				actionData.ExportMetrics(c.PromStorage)
+			}
+
+			actionDataJson, _ := json.Marshal(actionData)
+			if err := c.MQ.Publish("lagoon-actions", actionDataJson); err != nil {
+				opLog.Error(err, "error publishing message to mq")
+			}
 		}
 	}
 }