fix: restart workloads fix in app group (#5313)

Author: gireesh-naidu

go.mod

@@ -2,7 +2,7 @@ module github.com/devtron-labs/devtron
 
 go 1.21
 
-toolchain go1.21.8
+//toolchain go1.21.8
 
 require (
 	github.com/Masterminds/semver v1.5.0

pkg/deployment/manifest/deploymentTemplate/chartRef/ChartRefService.go

@@ -18,6 +18,7 @@ package chartRef
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -29,8 +30,10 @@ import (
 	util2 "github.com/devtron-labs/devtron/util"
 	dirCopy "github.com/otiai10/copy"
 	"go.uber.org/zap"
+	"golang.org/x/exp/maps"
 	"io/ioutil"
 	"k8s.io/helm/pkg/chartutil"
+	chart2 "k8s.io/helm/pkg/proto/hapi/chart"
 	"os"
 	"path"
 	"path/filepath"
@@ -70,19 +73,21 @@ type ChartRefFileOpService interface {
 	GetRefChart(chartRefId int) (string, string, string, string, error)
 	ExtractChartIfMissing(chartData []byte, refChartDir string, location string) (*bean.ChartDataInfo, error)
 	GetChartInBytes(chartRefId int, deleteChart bool) ([]byte, error)
-	GetChartBytesInBulk(chartRefIds []int, deleteChart bool) (map[int][]byte, error)
+	GetChartBytesForApps(ctx context.Context, appIdToAppName map[int]string) (map[int][]byte, error)
 }
 
 type ChartRefServiceImpl struct {
 	logger               *zap.SugaredLogger
 	chartRefRepository   chartRepoRepository.ChartRefRepository
 	chartTemplateService util.ChartTemplateService
 	mergeUtil            util.MergeUtil
+	chartRepository      chartRepoRepository.ChartRepository
 }
 
 func NewChartRefServiceImpl(logger *zap.SugaredLogger,
 	chartRefRepository chartRepoRepository.ChartRefRepository,
 	chartTemplateService util.ChartTemplateService,
+	chartRepository chartRepoRepository.ChartRepository,
 	mergeUtil util.MergeUtil) *ChartRefServiceImpl {
 	// cache devtron reference charts list
 	devtronChartList, _ := chartRefRepository.FetchAllNonUserUploadedChartInfo()
@@ -92,6 +97,7 @@ func NewChartRefServiceImpl(logger *zap.SugaredLogger,
 		chartRefRepository:   chartRefRepository,
 		chartTemplateService: chartTemplateService,
 		mergeUtil:            mergeUtil,
+		chartRepository:      chartRepository,
 	}
 }
 
@@ -309,23 +315,83 @@ func (impl *ChartRefServiceImpl) extractChartInBytes(chartRef *chartRepoReposito
 	}
 	return manifestByteArr, nil
 }
-func (impl *ChartRefServiceImpl) GetChartBytesInBulk(chartRefIds []int, performCleanup bool) (map[int][]byte, error) {
+
+func (impl *ChartRefServiceImpl) getChartPath(chartRef *chartRepoRepository.ChartRef) (string, error) {
+	refChartPath := filepath.Join(bean.RefChartDirPath, chartRef.Location)
+	// For user uploaded charts ChartData will be retrieved from DB
+	if chartRef.ChartData != nil {
+		chartInfo, err := impl.ExtractChartIfMissing(chartRef.ChartData, bean.RefChartDirPath, chartRef.Location)
+		if chartInfo != nil && chartInfo.TemporaryFolder != "" {
+			err1 := os.RemoveAll(chartInfo.TemporaryFolder)
+			if err1 != nil {
+				impl.logger.Errorw("error in deleting temp dir ", "err", err)
+			}
+		}
+	} else {
+		// For Devtron reference charts the chart will be load from the directory location
+	}
+	return refChartPath, nil
+}
+
+func (impl *ChartRefServiceImpl) GetChartBytesForApps(ctx context.Context, appIdToAppName map[int]string) (map[int][]byte, error) {
+
+	appIds := maps.Keys(appIdToAppName)
+	charts, err := impl.chartRepository.FindLatestChartByAppIds(appIds)
+	if err != nil {
+		impl.logger.Errorw("error in fetching chart", "err", err, "appIds", appIds)
+		return nil, err
+	}
+
+	chartRefIdTOAppIds := make(map[int][]int)
+	var chartRefIds []int
+	chartRefToChartVersion := make(map[int]string)
+
+	for _, chart := range charts {
+		chartRefIds = append(chartRefIds, chart.ChartRefId)
+		chartRefToChartVersion[chart.ChartRefId] = chart.ChartVersion
+		refAppIds, ok := chartRefIdTOAppIds[chart.ChartRefId]
+		if !ok {
+			refAppIds = make([]int, 0)
+		}
+		refAppIds = append(refAppIds, chart.AppId)
+		chartRefIdTOAppIds[chart.ChartRefId] = refAppIds
+	}
+
 	chartRefs, err := impl.chartRefRepository.FindByIds(chartRefIds)
 	if err != nil {
 		impl.logger.Errorw("error getting chart data", "chartRefIds", chartRefIds, "err", err)
 		return nil, err
 	}
-	chartRefIdToBytes := make(map[int][]byte)
+
+	appIdToBytes := make(map[int][]byte)
+
+	// this loops run with O(len(apps)) T.C
 	for _, chartRef := range chartRefs {
-		chartInBytes, err := impl.extractChartInBytes(chartRef, performCleanup)
+		refChartPath, err := impl.getChartPath(chartRef)
 		if err != nil {
 			impl.logger.Errorw("error in converting chart to bytes", "chartRefId", chartRef.Id, "err", err)
 			return nil, err
 		}
-		chartRefIdToBytes[chartRef.Id] = chartInBytes
+
+		refAppIds := chartRefIdTOAppIds[chartRef.Id]
+		for _, appId := range refAppIds {
+			chartMetaData := &chart2.Metadata{
+				Name:    appIdToAppName[appId],
+				Version: chartRefToChartVersion[chartRef.Id],
+			}
+			tempReferenceTemplateDir, err := impl.chartTemplateService.BuildChart(ctx, chartMetaData, refChartPath)
+			if err != nil {
+				impl.logger.Errorw("error in building chart", "chartMetaData", chartMetaData, "refChartPath", refChartPath)
+				return nil, err
+			}
+			chartInBytes, err := impl.chartTemplateService.LoadChartInBytes(tempReferenceTemplateDir, true)
+			appIdToBytes[appId] = chartInBytes
+		}
+
 	}
-	return chartRefIdToBytes, nil
+	return appIdToBytes, nil
 }
+
 func (impl *ChartRefServiceImpl) FetchCustomChartsInfo() ([]*bean.ChartDto, error) {
 	resultsMetadata, err := impl.chartRefRepository.GetAllChartMetadata()
 	if err != nil {

pkg/generateManifest/DeploymentTemplateService.go

@@ -46,6 +46,7 @@ import (
 	"github.com/gammazero/workerpool"
 	"github.com/go-pg/pg"
 	"go.uber.org/zap"
+	"golang.org/x/exp/maps"
 	"net/http"
 	"os"
 	"regexp"
@@ -422,6 +423,12 @@ func (impl DeploymentTemplateServiceImpl) GetRestartWorkloadData(ctx context.Con
 	var finalError error
 	for i, _ := range partitionedRequests {
 		req := partitionedRequests[i]
+		err = impl.setChartContent(ctx, req, appNameToId)
+		if err != nil {
+			impl.Logger.Errorw("error in setting chart content for apps", "appNames", maps.Keys(appNameToId), "err", err)
+			// continue processing next batch
+			continue
+		}
 		wp.Submit(func() {
 			resp, err := impl.helmAppClient.TemplateChartBulk(ctx, &gRPC.BulkInstallReleaseRequest{BulkInstallReleaseRequest: req})
 			if err != nil {

pkg/generateManifest/helper.go

@@ -17,12 +17,15 @@
 package generateManifest
 
 import (
+	"context"
 	"fmt"
 	"github.com/devtron-labs/common-lib/utils/yaml"
 	"github.com/devtron-labs/devtron/api/helm-app/bean"
 	"github.com/devtron-labs/devtron/api/helm-app/gRPC"
 	"github.com/devtron-labs/devtron/internal/sql/repository/app"
 	"github.com/devtron-labs/devtron/pkg/cluster/repository"
+	"go.opentelemetry.io/otel"
+	"golang.org/x/exp/maps"
 )
 
 func (impl DeploymentTemplateServiceImpl) constructRotatePodResponse(templateChartResponse []*gRPC.TemplateChartResponse, appNameToId map[string]int, environment *repository.Environment) (*RestartPodResponse, error) {
@@ -37,23 +40,15 @@ func (impl DeploymentTemplateServiceImpl) constructRotatePodResponse(templateCha
 		for _, manifest := range manifests {
 			gvk := manifest.GroupVersionKind()
 			name := manifest.GetName()
-			labels := manifest.GetLabels()
 			switch gvk.Kind {
 			case string(Deployment), string(StatefulSet), string(DemonSet), string(Rollout):
-				if labels != nil && labels[LabelReleaseKey] != "" {
-					resourceMeta = append(resourceMeta, &ResourceMetadata{
-						Name:             labels[LabelReleaseKey],
-						GroupVersionKind: gvk,
-					})
-				} else {
-					resourceMeta = append(resourceMeta, &ResourceMetadata{
-						Name:             name,
-						GroupVersionKind: gvk,
-					})
-				}
-
+				resourceMeta = append(resourceMeta, &ResourceMetadata{
+					Name:             name,
+					GroupVersionKind: gvk,
+				})
 			}
 		}
+
 		appIdToResourceIdentifier[appNameToId[tcResp.AppName]] = &ResourceIdentifierResponse{
 			ResourceMetaData: resourceMeta,
 			AppName:          appName,
@@ -91,10 +86,12 @@ func (impl DeploymentTemplateServiceImpl) constructInstallReleaseBulkReq(apps []
 	if appIdToInstallReleaseRequest == nil || len(appIdToInstallReleaseRequest) == 0 {
 		return nil, err
 	}
-	err = impl.setChartContent(applicationIds, appIdToInstallReleaseRequest)
-	if err != nil {
-		impl.Logger.Errorw("error in setting chart content", "appIds", applicationIds, "err", err)
-		return nil, err
+
+	for _, app := range apps {
+		if _, ok := appIdToInstallReleaseRequest[app.Id]; ok {
+			appIdToInstallReleaseRequest[app.Id].AppName = app.AppName
+			appIdToInstallReleaseRequest[app.Id].ChartRepository = ChartRepository
+		}
 	}
 
 	k8sServerVersion, err := impl.K8sUtil.GetKubeVersion()
@@ -111,8 +108,6 @@ func (impl DeploymentTemplateServiceImpl) constructInstallReleaseBulkReq(apps []
 		if _, ok := appIdToInstallReleaseRequest[app.Id]; ok {
 			appIdToInstallReleaseRequest[app.Id].ReleaseIdentifier = impl.getReleaseIdentifier(config, app, environment)
 			appIdToInstallReleaseRequest[app.Id].K8SVersion = k8sServerVersion.String()
-			appIdToInstallReleaseRequest[app.Id].AppName = app.AppName
-			appIdToInstallReleaseRequest[app.Id].ChartRepository = ChartRepository
 		}
 	}
 
@@ -122,30 +117,26 @@ func (impl DeploymentTemplateServiceImpl) constructInstallReleaseBulkReq(apps []
 	return installReleaseRequests, nil
 }
 
-func (impl DeploymentTemplateServiceImpl) setChartContent(appIds []int, appIdToInstallReleaseRequest map[int]*gRPC.InstallReleaseRequest) error {
-	charts, err := impl.chartRepository.FindLatestChartByAppIds(appIds)
-	if err != nil {
-		impl.Logger.Errorw("error in fetching chart", "err", err, "appIds", appIds)
-		return err
-	}
-	appIdToChartRefId := make(map[int]int)
-	var chartRefIds []int
-	for _, chart := range charts {
-		appIdToChartRefId[chart.AppId] = chart.ChartRefId
-		chartRefIds = append(chartRefIds, chart.ChartRefId)
+func (impl DeploymentTemplateServiceImpl) setChartContent(ctx context.Context, installReleaseRequests []*gRPC.InstallReleaseRequest, appNameToId map[string]int) error {
+	appIdToInstallReleaseRequest := make(map[int]*gRPC.InstallReleaseRequest)
+	requestAppNameToId := make(map[int]string)
+	for _, req := range installReleaseRequests {
+		appId := appNameToId[req.AppName]
+		requestAppNameToId[appId] = req.AppName
+		appIdToInstallReleaseRequest[appId] = req
 	}
-	chartRefIdToBytes, err := impl.chartRefService.GetChartBytesInBulk(chartRefIds, false)
+
+	_, span := otel.Tracer("orchestrator").Start(ctx, "chartRefService.GetChartBytesForApps")
+	appIdToBytes, err := impl.chartRefService.GetChartBytesForApps(ctx, requestAppNameToId)
+	span.End()
 	if err != nil {
-		impl.Logger.Errorw("error in fetching chartRefBean", "err", err, "chartRefIds", chartRefIds)
+		impl.Logger.Errorw("error in fetching chartRefBean", "err", err, "appNames", maps.Keys(appNameToId))
 		return err
 	}
-	for appId, chartRefId := range appIdToChartRefId {
-		if bytes, ok := chartRefIdToBytes[chartRefId]; ok {
-			if _, ok := appIdToInstallReleaseRequest[appId]; ok {
-				appIdToInstallReleaseRequest[appId].ChartContent = &gRPC.ChartContent{Content: bytes}
-			}
-		}
+	for appId, _ := range appIdToInstallReleaseRequest {
+		appIdToInstallReleaseRequest[appId].ChartContent = &gRPC.ChartContent{Content: appIdToBytes[appId]}
 	}
+
 	return err
 }