Merge pull request #33 from kuzxnia/fix_conn_leak

Fix conn leak

cli/cli.go

@@ -168,6 +168,8 @@ func provideWorkloadCommands() []*cobra.Command {
 		Use:     CommandConfigWorkload,
 		Short:   "Config",
 		GroupID: WorkloadGroup.ID,
+		PersistentPreRunE: persistentPreRunE,
+		PersistentPostRun: persistentPostRun,
 		RunE: func(cmd *cobra.Command, args []string) (err error) {
 			flags := cmd.Flags()
 			configFile, _ := flags.GetString(ConfigFile)
@@ -340,6 +342,15 @@ func provideOrchiestrationCommands() []*cobra.Command {
 			helmSet, _ := flags.GetStringSlice(FlagHelmSet)
 			workloadConfigPath, _ := flags.GetString(FlagWorkloadConfig)
 
+			cfg, err := ParseConfigFile(workloadConfigPath, false)
+			if err != nil {
+				return err
+			}
+			configValues, err := cfg.Values()
+			if err != nil {
+				return err
+			}
+
 			rsm := resourcemanager.ResourceManagerConfig{
 				KubeconfigPath: srcKubeconfigPath,
 				Context:        srcContext,
@@ -351,7 +362,7 @@ func provideOrchiestrationCommands() []*cobra.Command {
 				ResourceManagerConfig: rsm,
 				Name:                  args[0],
 				HelmValues:            helmSet,
-				WorkloadConfigString:  workloadConfigPath,
+				WorkloadConfigString:  configValues,
 			}
 
 			return UpgradeResources(&request)

cli/workload/config.go

@@ -17,7 +17,7 @@ import (
 
 // checks if process is running in local system
 // here should be cli config request - not lbot one
-func SetConfigWorkload(conn *grpc.ClientConn, parsedConfig *lbot.ConfigRequest) (err error) {
+func SetConfigWorkload(conn grpc.ClientConnInterface, parsedConfig *lbot.ConfigRequest) (err error) {
 	requestConfig := BuildConfigRequest(parsedConfig)
 
 	fmt.Println("🚀 Setting new config")

helm/workload/values.yaml

@@ -4,6 +4,6 @@ workload:
   name:
   replicas: 1
   agent:
-    image: kuzxnia/loadbot:v1.0.5
+    image: kuzxnia/loadbot:v1.0.6
     port: 1234
   config:

lbot/config.go

@@ -238,9 +238,9 @@ func (c *JobRequest) UnmarshalJSON(data []byte) (err error) {
 		Pace        uint64                 `json:"pace"`
 		DataSize    uint64                 `json:"data_size"`
 		BatchSize   uint64                 `json:"batch_size"`
-		Duration    config.Duration               `json:"duration"`
+		Duration    config.Duration        `json:"duration"`
 		Operations  uint64                 `json:"operations"`
-		Timeout     config.Duration               `json:"timeout"` // if not set, default
+		Timeout     config.Duration        `json:"timeout"` // if not set, default
 		Filter      map[string]interface{} `json:"filter"`
 	}
 	// default values
@@ -259,9 +259,9 @@ func (c *JobRequest) UnmarshalJSON(data []byte) (err error) {
 	c.Pace = tmp.Pace
 	c.DataSize = tmp.DataSize
 	c.BatchSize = tmp.BatchSize
-  c.Duration = tmp.Duration.Duration
+	c.Duration = tmp.Duration.Duration
 	c.Operations = tmp.Operations
-  c.Timeout = tmp.Timeout.Duration
+	c.Timeout = tmp.Timeout.Duration
 	c.Filter = tmp.Filter
 
 	return

lbot/database/client.go

@@ -36,7 +36,8 @@ func NewMongoClient(connectionString string, cfg *config.Job, schema *config.Sch
 	opts = opts.
 		ApplyURI(connectionString).
 		SetReadPreference(readpref.SecondaryPreferred()).
-		SetMaxConnecting(100).
+		SetMaxPoolSize(cfg.Connections * 2).
+		// SetMaxConnecting(100).
 		SetMaxConnIdleTime(90 * time.Second).
 		SetTimeout(cfg.Timeout)
 

lbot/lbot.go

@@ -40,6 +40,7 @@ func (l *Lbot) Run() error {
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 
 	for _, job := range l.Config.Jobs {
 		err = client.RunJob(*job)
@@ -56,6 +57,7 @@ func (l *Lbot) SetCommandState(command *database.Command, state database.Command
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 	command.State = state.String()
 	return client.SaveCommand(command)
 }
@@ -65,6 +67,7 @@ func (l *Lbot) SetWorkloadState(workload *database.Workload, state database.Work
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 	workload.State = state.String()
 	return client.SaveWorkload(workload)
 }
@@ -147,6 +150,7 @@ func (l *Lbot) InitAgent(id primitive.ObjectID, name string) error {
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 	ct, err := client.ClusterTime()
 	if err != nil {
 		return errors.Wrap(err, "get cluster time")
@@ -167,6 +171,7 @@ func (l *Lbot) AgentHeartBeat(id primitive.ObjectID, name string) error {
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 
 	agentStatus := database.AgentStatus{
 		Id:   id,
@@ -180,9 +185,11 @@ func (l *Lbot) HandleWorkload() {
 	log.Println("Fetching new workloads")
 	// todo: change to generic abstraction
 	client, err := database.NewInternalMongoClient(l.Config.ConnectionString)
+  defer client.Disconnect()
 	if err != nil {
 		return
 	}
+  defer client.Disconnect()
 
 	// todo: change to commands
 	workload, err := client.GetNewWorkloads()
@@ -218,6 +225,7 @@ func (l *Lbot) HandleCommand() {
 	if err != nil {
 		return
 	}
+  defer client.Disconnect()
 
 	log.Println("Fetching not finished commands")
 	// todo: change to commands
@@ -265,6 +273,7 @@ func (l *Lbot) IsMasterAgent(agentId primitive.ObjectID) (bool, error) {
 	if err != nil {
 		return false, err
 	}
+  defer client.Disconnect()
 	return client.IsMasterAgent(agentId)
 }
 
@@ -274,6 +283,7 @@ func (l *Lbot) UpdateRunningAgents() error {
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 
 	runningAgents, err := client.GetAgentWithHeartbeatWithin()
 	if err != nil {
@@ -297,6 +307,7 @@ func (l *Lbot) GetNextUnFinishedCommand() (*database.Command, error) {
 	if err != nil {
 		return nil, err
 	}
+  defer client.Disconnect()
 	return client.GetNextUnFinishedCommand()
 }
 
@@ -305,6 +316,7 @@ func (l *Lbot) AreWorkloadsFinished(command *database.Command) (bool, error) {
 	if err != nil {
 		return false, err
 	}
+  defer client.Disconnect()
 	workloads, err := client.GetCommandWorkloads(command)
 
 	finished := lo.Filter(workloads, func(w *database.Workload, index int) bool {
@@ -319,6 +331,7 @@ func (l *Lbot) GenerateWorkload(command *database.Command) ([]*database.Workload
 	if err != nil {
 		return nil, err
 	}
+  defer client.Disconnect()
 	ct, err := client.ClusterTime()
 	if err != nil {
 		return nil, errors.Wrap(err, "get cluster time")
@@ -356,6 +369,7 @@ func (l *Lbot) SaveWorkloads(ws []*database.Workload) error {
 	if err != nil {
 		return err
 	}
+  defer client.Disconnect()
 
 	var worklods []interface{}
 	for _, w := range ws {

lbot/resourcemanager/helm.go

@@ -61,7 +61,7 @@ func (c *HelmManager) Install(request *InstallRequest) (err error) {
 	installer.Labels = map[string]string{"role": "workload"}
 
 	options := values.Options{
-		Values:        append([]string{"workload.name=" + request.Name}, request.HelmValues...),
+		Values:        append([]string{"workload.name=" + request.Name, "workload.namespace=" + request.Namespace}, request.HelmValues...),
 		LiteralValues: []string{"workload.config=" + request.WorkloadConfigString},
 	}
 
@@ -106,7 +106,7 @@ func (c *HelmManager) Upgrade(request *UpgradeRequest) (err error) {
 	upgrader.Labels = map[string]string{"role": "workload"}
 
 	options := values.Options{
-		Values:        append([]string{"workload.name=" + request.Name}, request.HelmValues...),
+		Values:        append([]string{"workload.name=" + request.Name, "workload.namespace=" + request.Namespace}, request.HelmValues...),
 		LiteralValues: []string{"workload.config=" + request.WorkloadConfigString},
 	}
 

lbot/schema/data_provider.go

@@ -30,7 +30,7 @@ type LiveDataProvider struct {
 func NewLiveDataProvider(job *config.Job, schema *config.Schema) *LiveDataProvider {
 	return &LiveDataProvider{
 		job:           job,
-		dataGenerator: NewDataGenerator(schema, job.DataSize),
+		dataGenerator: NewDataGenerator(schema, int(job.DataSize)),
 	}
 }
 

lbot/schema/generator.go

@@ -3,6 +3,8 @@ package schema
 import (
 	"errors"
 	"math/rand"
+	"time"
+	"unsafe"
 
 	"github.com/go-faker/faker/v4"
 	"github.com/go-faker/faker/v4/pkg/options"
@@ -63,8 +65,8 @@ type DataGenerator interface {
 	GenerateFromTemplate(interface{}) (interface{}, error)
 }
 
-func NewDataGenerator(schema *config.Schema, dataSize uint64) DataGenerator {
-  // todo: check size of object using, unsafe.Sizeof( )
+func NewDataGenerator(schema *config.Schema, dataSize int) DataGenerator {
+	// todo: check size of object using, unsafe.Sizeof( )
 
 	if schema != nil {
 		return DataGenerator(
@@ -82,7 +84,7 @@ func NewDataGenerator(schema *config.Schema, dataSize uint64) DataGenerator {
 }
 
 type MeasurableDataGenerator struct {
-	dataSize uint64
+	dataSize int
 }
 
 func (g *MeasurableDataGenerator) Generate() (interface{}, error) {
@@ -109,14 +111,31 @@ func (g *MeasurableDataGenerator) GenerateFromTemplate(template interface{}) (in
 	}
 }
 
-const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+const (
+	letterBytes   = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	letterIdxBits = 6                    // 6 bits to represent a letter index
+	letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
+	letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
+)
+
+var src = rand.NewSource(time.Now().UnixNano())
 
-func randStringBytes(n uint64) string {
+func randStringBytes(n int) string {
 	b := make([]byte, n)
-	for i := range b {
-		b[i] = letterBytes[rand.Intn(len(letterBytes))]
+	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
+	for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
+		if remain == 0 {
+			cache, remain = src.Int63(), letterIdxMax
+		}
+		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
+			b[i] = letterBytes[idx]
+			i--
+		}
+		cache >>= letterIdxBits
+		remain--
 	}
-	return string(b)
+
+	return *(*string)(unsafe.Pointer(&b))
 }
 
 type StructuralizableDataGenerator struct {