fix: use sync map for lazy load

Author: Yifshao

domain/project.go

@@ -3,6 +3,7 @@ package domain
 import (
 	"fmt"
 	"strconv"
+	"sync"
 
 	"github.com/aliyun/aliyun-pai-featurestore-go-sdk/v2/api"
 	"github.com/aliyun/aliyun-pai-featurestore-go-sdk/v2/constants"
@@ -15,21 +16,18 @@ import (
 type Project struct {
 	*api.Project
 	OnlineStore      OnlineStore
-	FeatureViewMap   map[string]FeatureView
+	FeatureViewMap   sync.Map
 	FeatureEntityMap map[string]*FeatureEntity
-	ModelMap         map[string]*Model
-	LabelTableMap    map[int]*LabelTable
+	ModelMap         sync.Map
+	LabelTableMap    sync.Map
 
 	apiClient *api.APIClient
 }
 
 func NewProject(p *api.Project, isInitClient bool) *Project {
 	project := Project{
 		Project:          p,
-		FeatureViewMap:   make(map[string]FeatureView),
 		FeatureEntityMap: make(map[string]*FeatureEntity),
-		ModelMap:         make(map[string]*Model),
-		LabelTableMap:    make(map[int]*LabelTable),
 	}
 
 	switch p.OnlineDatasourceType {
@@ -89,50 +87,58 @@ func (p *Project) SetApiClient(apiClient *api.APIClient) {
 }
 
 func (p *Project) GetFeatureView(name string) FeatureView {
-	if _, exists := p.FeatureViewMap[name]; !exists {
-		err := p.loadFeatureView(name)
-		if err != nil {
-			return nil
-		}
+	if value, exists := p.FeatureViewMap.Load(name); exists {
+		return value.(FeatureView)
 	}
-
-	return p.FeatureViewMap[name]
+	if err := p.loadFeatureView(name); err != nil {
+		return nil
+	}
+	if value, exists := p.FeatureViewMap.Load(name); exists {
+		return value.(FeatureView)
+	}
+	return nil
 }
 
 func (p *Project) GetFeatureEntity(name string) *FeatureEntity {
 	return p.FeatureEntityMap[name]
 }
 
 func (p *Project) GetLabelTable(labelTableId int) *LabelTable {
-	if _, exists := p.LabelTableMap[labelTableId]; !exists {
-		err := p.loadLabelTable(labelTableId)
-		if err != nil {
-			return nil
-		}
+	if value, exists := p.LabelTableMap.Load(labelTableId); exists {
+		return value.(*LabelTable)
 	}
-
-	return p.LabelTableMap[labelTableId]
+	if err := p.loadLabelTable(labelTableId); err != nil {
+		return nil
+	}
+	if value, exists := p.LabelTableMap.Load(labelTableId); exists {
+		return value.(*LabelTable)
+	}
+	return nil
 }
 
 func (p *Project) GetModel(name string) *Model {
-	if _, exists := p.ModelMap[name]; !exists {
-		err := p.loadModelFeature(name)
-		if err != nil {
-			return nil
-		}
+	if value, exists := p.ModelMap.Load(name); exists {
+		return value.(*Model)
 	}
-
-	return p.ModelMap[name]
+	if err := p.loadModelFeature(name); err != nil {
+		return nil
+	}
+	if value, exists := p.ModelMap.Load(name); exists {
+		return value.(*Model)
+	}
+	return nil
 }
 func (p *Project) GetModelFeature(name string) *Model {
-	if _, exists := p.ModelMap[name]; !exists {
-		err := p.loadModelFeature(name)
-		if err != nil {
-			return nil
-		}
+	if value, exists := p.ModelMap.Load(name); exists {
+		return value.(*Model)
 	}
-
-	return p.ModelMap[name]
+	if err := p.loadModelFeature(name); err != nil {
+		return nil
+	}
+	if value, exists := p.ModelMap.Load(name); exists {
+		return value.(*Model)
+	}
+	return nil
 }
 
 func (p *Project) loadFeatureView(featureViewName string) error {
@@ -166,7 +172,7 @@ func (p *Project) loadFeatureView(featureViewName string) error {
 				return fmt.Errorf("feature entity not exist, name=%s", featureView.FeatureEntityName)
 			}
 			featureViewDomain := NewFeatureView(featureView, p, entity)
-			p.FeatureViewMap[featureView.Name] = featureViewDomain
+			p.FeatureViewMap.Store(featureView.Name, featureViewDomain)
 		}
 
 		if len(listFeatureViews.FeatureViews) == 0 || pageSize*pageNumber > listFeatureViews.TotalCount {
@@ -186,7 +192,7 @@ func (p *Project) loadLabelTable(labelTableId int) error {
 		return err
 	}
 	labelTableDomain := NewLabelTable(getLabelTableResponse.LabelTable)
-	p.LabelTableMap[labelTableId] = labelTableDomain
+	p.LabelTableMap.Store(labelTableId, labelTableDomain)
 
 	return nil
 }
@@ -207,17 +213,13 @@ func (p *Project) loadModelFeature(modelFeatureName string) error {
 				return err
 			}
 			model := getModelFeatureResponse.Model
-			var labelTableDomain *LabelTable
-			if _, exists := p.LabelTableMap[model.LabelTableId]; !exists {
-				err := p.loadLabelTable(model.LabelTableId)
-				if err != nil {
-					fmt.Printf("get label table error, err=%v", err)
-					return err
-				}
+			labelTableDomain := p.GetLabelTable(model.LabelTableId)
+			if labelTableDomain == nil {
+				fmt.Printf("label table not exist, id=%d", model.LabelTableId)
+				return fmt.Errorf("label table not exist, id=%d", model.LabelTableId)
 			}
-			labelTableDomain = p.LabelTableMap[model.LabelTableId]
 			modelDomain := NewModel(model, p, labelTableDomain)
-			p.ModelMap[model.Name] = modelDomain
+			p.ModelMap.Store(model.Name, modelDomain)
 		}
 
 		if len(listModelFeatures.Models) == 0 || pageSize*pageNumber > listModelFeatures.TotalCount {

featurestore/feature_store_client.go

@@ -266,6 +266,7 @@ func (c *FeatureStoreClient) LoadProjectData() error {
 		p.Signature = c.signature
 
 		project := domain.NewProject(p, c.datasourceInitClient)
+		project.SetApiClient(c.client)
 		projectData[project.ProjectName] = project
 
 		var (
@@ -319,7 +320,7 @@ func (c *FeatureStoreClient) LoadProjectData() error {
 				}
 
 				featureViewDomain := domain.NewFeatureView(featureView, project, project.FeatureEntityMap[featureView.FeatureEntityName])
-				project.FeatureViewMap[featureView.Name] = featureViewDomain
+				project.FeatureViewMap.Store(featureView.Name, featureViewDomain)
 
 			}
 
@@ -333,7 +334,6 @@ func (c *FeatureStoreClient) LoadProjectData() error {
 
 		pagenumber = 1
 		// get model
-		labelTableMap := make(map[string]*domain.LabelTable)
 		for {
 			listModelsResponse, err := c.client.FsModelApi.ListModels(pagesize, pagenumber, strconv.Itoa(project.ProjectId))
 			if err != nil {
@@ -348,20 +348,13 @@ func (c *FeatureStoreClient) LoadProjectData() error {
 					return err
 				}
 				model := getModelResponse.Model
-				var labelTableDomain *domain.LabelTable
-				if labelTable, exists := labelTableMap[model.LabelDatasourceTable]; !exists || labelTable == nil {
-					getLabelTableResponse, err := c.client.LabelTableApi.GetLabelTableByID(strconv.Itoa(model.LabelTableId))
-					if err != nil {
-						c.logError(fmt.Errorf("get label table error, err=%v", err))
-						return err
-					}
-					labelTableDomain = domain.NewLabelTable(getLabelTableResponse.LabelTable)
-					labelTableMap[model.LabelDatasourceTable] = labelTableDomain
-				} else {
-					labelTableDomain = labelTable
+				labelTable := project.GetLabelTable(model.LabelTableId)
+				if labelTable == nil {
+					c.logError(fmt.Errorf("not found label table, labelTableId:%d", model.LabelTableId))
+					return fmt.Errorf("not found label table, labelTableId:%d", model.LabelTableId)
 				}
-				modelDomain := domain.NewModel(model, project, labelTableDomain)
-				project.ModelMap[model.Name] = modelDomain
+				modelDomain := domain.NewModel(model, project, labelTable)
+				project.ModelMap.Store(model.Name, modelDomain)
 
 			}
 
@@ -472,6 +465,16 @@ func (c *FeatureStoreClient) lazyLoadProjectData() error {
 }
 
 func (c *FeatureStoreClient) loopLoadProjectData() {
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				fmt.Println("Recovered from panic:", r)
+			}
+		}()
+
+		c.LoadProjectData()
+	}()
+
 	ticker := time.NewTicker(time.Minute)
 	defer ticker.Stop()
 	for {