Feat: Cloud Virtual Machine Billing

controllers/account/controllers/account_controller.go

@@ -22,10 +22,17 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"os"
 	"strconv"
 	"strings"
 	"time"
 
+	"go.mongodb.org/mongo-driver/bson/primitive"
+
+	"github.com/google/uuid"
+
+	gonanoid "github.com/matoous/go-nanoid/v2"
+
 	"gorm.io/gorm"
 
 	"k8s.io/client-go/rest"
@@ -74,6 +81,7 @@ type AccountReconciler struct {
 	Logger                 logr.Logger
 	AccountSystemNamespace string
 	DBClient               database.Account
+	CVMDBClient            database.CVM
 	MongoDBURI             string
 	Activities             pkgtypes.Activities
 	DefaultDiscount        pkgtypes.RechargeDiscount
@@ -399,3 +407,68 @@ func getAmountWithDiscount(amount int64, discount pkgtypes.RechargeDiscount) int
 	}
 	return int64(math.Ceil(float64(amount) * r / 100))
 }
+
+func (r *AccountReconciler) BillingCVM() error {
+	cvmMap, err := r.CVMDBClient.GetPendingStateInstance(os.Getenv("LOCAL_REGION"))
+	if err != nil {
+		return fmt.Errorf("get pending state instance failed: %v", err)
+	}
+	for userInfo, cvms := range cvmMap {
+		fmt.Println("billing cvm", userInfo, cvms)
+		userUID, namespace := strings.Split(userInfo, "/")[0], strings.Split(userInfo, "/")[1]
+		appCosts := make([]resources.AppCost, len(cvms))
+		cvmTotalAmount := 0.0
+		cvmIDs := make([]primitive.ObjectID, len(cvms))
+		cvmIDsDetail := make([]string, 0)
+		for i := range cvms {
+			appCosts[i] = resources.AppCost{
+				Amount: int64(cvms[i].Amount * BaseUnit),
+				Name:   cvms[i].InstanceName,
+			}
+			cvmTotalAmount += cvms[i].Amount
+			cvmIDs[i] = cvms[i].ID
+			cvmIDsDetail = append(cvmIDsDetail, cvms[i].ID.String())
+		}
+		userQueryOpts := pkgtypes.UserQueryOpts{UID: uuid.MustParse(userUID)}
+		user, err := r.AccountV2.GetUserCr(&userQueryOpts)
+		if err != nil {
+			if err == gorm.ErrRecordNotFound {
+				fmt.Println("user not found", userQueryOpts)
+				continue
+			}
+			return fmt.Errorf("get user failed: %v", err)
+		}
+		id, err := gonanoid.New(12)
+		if err != nil {
+			return fmt.Errorf("generate billing id error: %v", err)
+		}
+		billing := &resources.Billing{
+			OrderID:   id,
+			AppCosts:  appCosts,
+			Type:      accountv1.Consumption,
+			Namespace: namespace,
+			AppType:   resources.AppType[resources.CVM],
+			Amount:    int64(cvmTotalAmount * BaseUnit),
+			Owner:     user.CrName,
+			Time:      time.Now().UTC(),
+			Status:    resources.Settled,
+			Detail:    "{" + strings.Join(cvmIDsDetail, ",") + "}",
+		}
+		err = r.AccountV2.AddDeductionBalanceWithFunc(&pkgtypes.UserQueryOpts{UID: user.UserUID}, billing.Amount, func() error {
+			if saveErr := r.DBClient.SaveBillings(billing); saveErr != nil {
+				return fmt.Errorf("save billing failed: %v", saveErr)
+			}
+			return nil
+		}, func() error {
+			if saveErr := r.CVMDBClient.SetDoneStateInstance(cvmIDs...); saveErr != nil {
+				return fmt.Errorf("set done state instance failed: %v", saveErr)
+			}
+			return nil
+		})
+		if err != nil {
+			return fmt.Errorf("add balance failed: %v", err)
+		}
+		fmt.Printf("billing cvm success %#+v\n", billing)
+	}
+	return nil
+}

controllers/account/controllers/account_controller_test.go

@@ -17,8 +17,16 @@ limitations under the License.
 package controllers
 
 import (
+	"context"
+	"os"
 	"testing"
 
+	ctrl "sigs.k8s.io/controller-runtime"
+
+	"github.com/labring/sealos/controllers/pkg/database"
+	"github.com/labring/sealos/controllers/pkg/database/cockroach"
+	"github.com/labring/sealos/controllers/pkg/database/mongo"
+
 	"github.com/labring/sealos/controllers/pkg/types"
 )
 
@@ -82,3 +90,49 @@ func Test_getAmountWithDiscount(t *testing.T) {
 		})
 	}
 }
+
+func TestAccountReconciler_BillingCVM(t *testing.T) {
+	dbCtx := context.Background()
+	cvmDBClient, err := mongo.NewMongoInterface(dbCtx, os.Getenv(database.CVMMongoURI))
+	if err != nil {
+		t.Fatalf("unable to connect to mongo: %v", err)
+	}
+	defer func() {
+		if cvmDBClient != nil {
+			err := cvmDBClient.Disconnect(dbCtx)
+			if err != nil {
+				t.Errorf("unable to disconnect from mongo: %v", err)
+			}
+		}
+	}()
+	v2Account, err := cockroach.NewCockRoach(os.Getenv(database.GlobalCockroachURI), os.Getenv(database.LocalCockroachURI))
+	if err != nil {
+		t.Fatalf("unable to connect to cockroach: %v", err)
+	}
+	defer func() {
+		err := v2Account.Close()
+		if err != nil {
+			t.Errorf("unable to disconnect from cockroach: %v", err)
+		}
+	}()
+	DBClient, err := mongo.NewMongoInterface(dbCtx, os.Getenv(database.MongoURI))
+	if err != nil {
+		t.Fatalf("unable to connect to mongo: %v", err)
+	}
+	defer func() {
+		err := DBClient.Disconnect(dbCtx)
+		if err != nil {
+			t.Errorf("unable to disconnect from mongo: %v", err)
+		}
+	}()
+
+	r := &AccountReconciler{
+		AccountV2:   v2Account,
+		DBClient:    DBClient,
+		CVMDBClient: cvmDBClient,
+		Logger:      ctrl.Log.WithName("controllers").WithName("AccountReconciler"),
+	}
+	if err := r.BillingCVM(); err != nil {
+		t.Errorf("AccountReconciler.BillingCVM() error = %v", err)
+	}
+}

controllers/account/main.go

@@ -22,6 +22,8 @@ import (
 	"os"
 	"time"
 
+	"github.com/labring/sealos/controllers/pkg/utils/env"
+
 	"github.com/labring/sealos/controllers/pkg/types"
 
 	"github.com/labring/sealos/controllers/pkg/database/cockroach"
@@ -143,6 +145,23 @@ func main() {
 			setupLog.Error(err, "unable to disconnect from mongo")
 		}
 	}()
+	var cvmDBClient database.Interface
+	cvmURI := os.Getenv(database.CVMMongoURI)
+	if cvmURI != "" {
+		cvmDBClient, err = mongo.NewMongoInterface(dbCtx, cvmURI)
+		if err != nil {
+			setupLog.Error(err, "unable to connect to mongo")
+			os.Exit(1)
+		}
+	}
+	defer func() {
+		if cvmDBClient != nil {
+			err := cvmDBClient.Disconnect(dbCtx)
+			if err != nil {
+				setupLog.Error(err, "unable to disconnect from mongo")
+			}
+		}
+	}()
 	v2Account, err := cockroach.NewCockRoach(os.Getenv(database.GlobalCockroachURI), os.Getenv(database.LocalCockroachURI))
 	if err != nil {
 		setupLog.Error(err, "unable to connect to cockroach")
@@ -155,10 +174,11 @@ func main() {
 		}
 	}()
 	accountReconciler := &controllers.AccountReconciler{
-		Client:    mgr.GetClient(),
-		Scheme:    mgr.GetScheme(),
-		DBClient:  dbClient,
-		AccountV2: v2Account,
+		Client:      mgr.GetClient(),
+		Scheme:      mgr.GetScheme(),
+		DBClient:    dbClient,
+		AccountV2:   v2Account,
+		CVMDBClient: cvmDBClient,
 	}
 	billingInfoQueryReconciler := &controllers.BillingInfoQueryReconciler{
 		Client:     mgr.GetClient(),
@@ -269,6 +289,24 @@ func main() {
 		os.Exit(1)
 	}
 
+	go func() {
+		if cvmDBClient == nil {
+			setupLog.Info("CVM DB client is nil, skip billing cvm")
+			return
+		}
+		ticker := time.NewTicker(env.GetDurationEnvWithDefault("BILLING_CVM_INTERVAL", 10*time.Minute))
+		defer ticker.Stop()
+		for {
+			setupLog.Info("start billing cvm", "time", time.Now().Format(time.RFC3339))
+			err := accountReconciler.BillingCVM()
+			if err != nil {
+				setupLog.Error(err, "fail to billing cvm")
+			}
+			setupLog.Info("end billing cvm", "time", time.Now().Format(time.RFC3339))
+			<-ticker.C
+		}
+	}()
+
 	setupLog.Info("starting manager")
 	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
 		setupLog.Error(err, "fail to run manager")

controllers/pkg/database/cockroach/accountv2.go

@@ -95,10 +95,35 @@ func (c *Cockroach) CreateRegion(region *types.Region) error {
 	return nil
 }
 
-func (c *Cockroach) GetUserCr(ops *types.UserQueryOpts) (*types.RegionUserCr, error) {
+func (c *Cockroach) GetUser(ops *types.UserQueryOpts) (*types.User, error) {
 	if err := checkOps(ops); err != nil {
 		return nil, err
 	}
+	queryUser := &types.User{}
+	if ops.UID != uuid.Nil {
+		queryUser.UID = ops.UID
+	}
+	if ops.ID != "" {
+		queryUser.ID = ops.ID
+	}
+	var user types.User
+	if err := c.DB.Where(queryUser).First(&user).Error; err != nil {
+		return nil, fmt.Errorf("failed to get user: %v", err)
+	}
+	return &user, nil
+}
+
+func (c *Cockroach) GetUserCr(ops *types.UserQueryOpts) (*types.RegionUserCr, error) {
+	if ops.UID == uuid.Nil && ops.Owner == "" {
+		if ops.ID == "" {
+			return nil, fmt.Errorf("empty query opts")
+		}
+		user, err := c.GetUser(ops)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get user: %v", err)
+		}
+		ops.UID = user.UID
+	}
 	query := &types.RegionUserCr{
 		CrName: ops.Owner,
 	}
@@ -131,7 +156,7 @@ func (c *Cockroach) GetUserUID(ops *types.UserQueryOpts) (uuid.UUID, error) {
 }
 
 func checkOps(ops *types.UserQueryOpts) error {
-	if ops.Owner == "" && ops.UID == uuid.Nil {
+	if ops.Owner == "" && ops.UID == uuid.Nil && ops.ID == "" {
 		return fmt.Errorf("empty query opts")
 	}
 	return nil
@@ -304,6 +329,18 @@ func (c *Cockroach) AddDeductionBalance(ops *types.UserQueryOpts, amount int64)
 	})
 }
 
+func (c *Cockroach) AddDeductionBalanceWithFunc(ops *types.UserQueryOpts, amount int64, preDo, postDo func() error) error {
+	return c.DB.Transaction(func(tx *gorm.DB) error {
+		if err := preDo(); err != nil {
+			return err
+		}
+		if err := c.updateBalance(tx, ops, amount, true, true); err != nil {
+			return err
+		}
+		return postDo()
+	})
+}
+
 func (c *Cockroach) CreateAccount(ops *types.UserQueryOpts, account *types.Account) (*types.Account, error) {
 	if ops.UID == uuid.Nil {
 		user, err := c.GetUserCr(ops)

controllers/pkg/database/interface.go

@@ -18,6 +18,8 @@ import (
 	"context"
 	"time"
 
+	"go.mongodb.org/mongo-driver/bson/primitive"
+
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/labring/sealos/controllers/pkg/database/cockroach"
@@ -33,6 +35,12 @@ import (
 type Interface interface {
 	Account
 	Traffic
+	CVM
+}
+
+type CVM interface {
+	GetPendingStateInstance(regionUID string) (cvmMap map[string][]types.CVMBilling, err error)
+	SetDoneStateInstance(instanceIDs ...primitive.ObjectID) error
 }
 
 type Account interface {
@@ -81,6 +89,7 @@ type Traffic interface {
 type AccountV2 interface {
 	Close() error
 	GetUserCr(user *types.UserQueryOpts) (*types.RegionUserCr, error)
+	GetUser(ops *types.UserQueryOpts) (*types.User, error)
 	CreateUser(oAuth *types.OauthProvider, regionUserCr *types.RegionUserCr, user *types.User, workspace *types.Workspace, userWorkspace *types.UserWorkspace) error
 	GetAccount(user *types.UserQueryOpts) (*types.Account, error)
 	GetUserOauthProvider(ops *types.UserQueryOpts) (*types.OauthProvider, error)
@@ -97,6 +106,7 @@ type AccountV2 interface {
 	TransferAccountV1(owner string, account *types.Account) (*types.Account, error)
 	GetUserAccountRechargeDiscount(user *types.UserQueryOpts) (*types.RechargeDiscount, error)
 	AddDeductionBalance(user *types.UserQueryOpts, balance int64) error
+	AddDeductionBalanceWithFunc(ops *types.UserQueryOpts, amount int64, preDo, postDo func() error) error
 }
 
 type Creator interface {
@@ -118,6 +128,7 @@ type MeteringOwnerTimeResult struct {
 
 const (
 	MongoURI           = "MONGO_URI"
+	CVMMongoURI        = "CVM_MONGO_URI"
 	GlobalCockroachURI = "GLOBAL_COCKROACH_URI"
 	LocalCockroachURI  = "LOCAL_COCKROACH_URI"
 	TrafficMongoURI    = "TRAFFIC_MONGO_URI"