From 3a4a347455697ff15bcb1a1e5528daa2ebb81609 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E9=9D=99?= Date: Thu, 16 Oct 2025 16:57:45 +0800 Subject: [PATCH 1/5] add multi dimension value support for diversity_rule_sort --- recconf/recconf.go | 5 + sort/diversity_rule.go | 120 +++++++++++++ sort/diversity_rule_sort.go | 62 ++++--- sort/diversity_rule_sort_test.go | 283 ++++++++++++++++++++++++++++--- 4 files changed, 426 insertions(+), 44 deletions(-) diff --git a/recconf/recconf.go b/recconf/recconf.go index 6fa94a4..ac2bd21 100644 --- a/recconf/recconf.go +++ b/recconf/recconf.go @@ -765,8 +765,13 @@ type SortConfig struct { ExclusionRules []ExclusionRuleConfig TimeInterval int BoostScoreByWeightDao BoostScoreByWeightDaoConfig + MultiValueDimensionConf []MultiValueDimensionConfig } +type MultiValueDimensionConfig struct { + DimensionName string + Delimiter string +} type BoostScoreByWeightDaoConfig struct { DaoConfig ItemFieldName string diff --git a/sort/diversity_rule.go b/sort/diversity_rule.go index 7d59b66..d06d6aa 100644 --- a/sort/diversity_rule.go +++ b/sort/diversity_rule.go @@ -5,9 +5,18 @@ import ( "github.com/alibaba/pairec/v2/module" "github.com/alibaba/pairec/v2/recconf" + "github.com/alibaba/pairec/v2/utils" ) type DiversityMatchFunc func(item *module.Item) bool + +type DiversityRuleInterface interface { + Match(item *module.Item, itemList []*module.Item) bool + GetWeight() int +} + +var _ DiversityRuleInterface = (*DiversityRule)(nil) + type DiversityRule struct { DiversityRuleConfig recconf.DiversityRuleConfig DimensionItemMap map[module.ItemId]string @@ -86,3 +95,114 @@ func (r *DiversityRule) GetWeight() int { return r.DiversityRuleConfig.Weight } + +type DiversityRuleMultiDimension struct { + DiversityRuleConfig recconf.DiversityRuleConfig + DimensionItemMap map[module.ItemId][]any + multiDimensionMap map[int]recconf.MultiValueDimensionConfig +} + +func NewDiversityRuleMultiDimension(config recconf.DiversityRuleConfig, size int, multiDimensionMap map[int]recconf.MultiValueDimensionConfig) *DiversityRuleMultiDimension { + rule := DiversityRuleMultiDimension{ + DiversityRuleConfig: config, + DimensionItemMap: make(map[module.ItemId][]any, size), + multiDimensionMap: multiDimensionMap, + } + + return &rule +} + +func (r *DiversityRuleMultiDimension) GetDimensionValue(item *module.Item) []any { + if value, ok := r.DimensionItemMap[item.Id]; ok { + return value + } + + var dimensionValues []any + for i, dimension := range r.DiversityRuleConfig.Dimensions { + multiDimensionConf, ok := r.multiDimensionMap[i] + if ok { + value := item.StringProperty(dimension) + strs := strings.Split(value, multiDimensionConf.Delimiter) + dimensionValues = append(dimensionValues, strs) + } else { + value := item.StringProperty(dimension) + dimensionValues = append(dimensionValues, value) + } + } + + r.DimensionItemMap[item.Id] = dimensionValues + + return r.DimensionItemMap[item.Id] +} + +func (r *DiversityRuleMultiDimension) Match(item *module.Item, itemList []*module.Item) bool { + size := len(itemList) + + itemDimensionValues := r.GetDimensionValue(item) + if r.DiversityRuleConfig.IntervalSize > 0 && size >= r.DiversityRuleConfig.IntervalSize { + end := size + begin := size - r.DiversityRuleConfig.IntervalSize + sameValue := 1 + for i := end - 1; i >= begin; i-- { + if r.isDimensionValuesEqual(itemDimensionValues, r.GetDimensionValue(itemList[i])) { + sameValue++ + } else { + break + } + } + + if sameValue > r.DiversityRuleConfig.IntervalSize { + return false + } + + } + if r.DiversityRuleConfig.WindowSize > 0 && + r.DiversityRuleConfig.FrequencySize > 0 && + r.DiversityRuleConfig.WindowSize > r.DiversityRuleConfig.FrequencySize { + end := size + begin := size - r.DiversityRuleConfig.WindowSize + 1 + if begin < 0 { + begin = 0 + } + + sameValue := 1 + for i := begin; i < end; i++ { + if r.isDimensionValuesEqual(itemDimensionValues, r.GetDimensionValue(itemList[i])) { + sameValue++ + } + + if sameValue > r.DiversityRuleConfig.FrequencySize { + return false + } + } + } + return true +} + +func (r *DiversityRuleMultiDimension) GetWeight() int { + return r.DiversityRuleConfig.Weight +} +func (r *DiversityRuleMultiDimension) isDimensionValuesEqual(left, right []any) bool { + if len(left) != len(right) { + return false + } + for i := 0; i < len(left); i++ { + switch left[i].(type) { + case string: + if left[i] != right[i] { + return false + } + case []string: + leftValues := left[i].([]string) + if rightValues, ok := right[i].([]string); ok { + if !utils.StringContains(leftValues, rightValues) { + return false + } + + } else { + return false + } + } + } + return true +} diff --git a/sort/diversity_rule_sort.go b/sort/diversity_rule_sort.go index 915574e..8934055 100644 --- a/sort/diversity_rule_sort.go +++ b/sort/diversity_rule_sort.go @@ -10,26 +10,29 @@ import ( ) type DiversityRuleSort struct { - diversitySize int - diversityRules []recconf.DiversityRuleConfig - exclusionRules []recconf.ExclusionRuleConfig - excludeRecallMap map[string]bool - filterParam *module.FilterParam - cloneInstances map[string]*DiversityRuleSort - name string - exploreItemSize int + diversitySize int + diversityRules []recconf.DiversityRuleConfig + exclusionRules []recconf.ExclusionRuleConfig + excludeRecallMap map[string]bool + filterParam *module.FilterParam + cloneInstances map[string]*DiversityRuleSort + name string + exploreItemSize int + multiValueDimensionConf []recconf.MultiValueDimensionConfig + multiDimensionMaps []map[int]recconf.MultiValueDimensionConfig // size is equal config.DiversityRules, per entry is map of multi dimension config for each diversity rule of dimensions } func NewDiversityRuleSort(config recconf.SortConfig) *DiversityRuleSort { sort := DiversityRuleSort{ - diversitySize: config.DiversitySize, - diversityRules: config.DiversityRules, - exclusionRules: config.ExclusionRules, - excludeRecallMap: make(map[string]bool, len(config.ExcludeRecalls)), - filterParam: nil, - name: config.Name, - cloneInstances: make(map[string]*DiversityRuleSort), - exploreItemSize: -1, + diversitySize: config.DiversitySize, + diversityRules: config.DiversityRules, + exclusionRules: config.ExclusionRules, + excludeRecallMap: make(map[string]bool, len(config.ExcludeRecalls)), + filterParam: nil, + name: config.Name, + cloneInstances: make(map[string]*DiversityRuleSort), + exploreItemSize: -1, + multiValueDimensionConf: config.MultiValueDimensionConf, } for _, recallName := range config.ExcludeRecalls { @@ -43,6 +46,22 @@ func NewDiversityRuleSort(config recconf.SortConfig) *DiversityRuleSort { if config.ExploreItemSize > 0 { sort.exploreItemSize = config.ExploreItemSize } + if len(config.MultiValueDimensionConf) > 0 { + for _, diversityRuleConfig := range sort.diversityRules { + multiDimensionMap := make(map[int]recconf.MultiValueDimensionConfig) + for i, dimension := range diversityRuleConfig.Dimensions { + for _, multiDimension := range config.MultiValueDimensionConf { + if multiDimension.DimensionName == dimension { + multiDimensionMap[i] = multiDimension + break + } + } + } + sort.multiDimensionMaps = append(sort.multiDimensionMaps, multiDimensionMap) + + } + + } return &sort } @@ -68,9 +87,14 @@ func (s *DiversityRuleSort) Sort(sortData *SortData) error { return nil } -func (s *DiversityRuleSort) createDiversityRules(size int) (ret []*DiversityRule) { - for _, config := range s.diversityRules { - rule := NewDiversityRule(config, size) +func (s *DiversityRuleSort) createDiversityRules(size int) (ret []DiversityRuleInterface) { + for i, config := range s.diversityRules { + var rule DiversityRuleInterface + if len(s.multiValueDimensionConf) > 0 { + rule = NewDiversityRuleMultiDimension(config, size, s.multiDimensionMaps[i]) + } else { + rule = NewDiversityRule(config, size) + } ret = append(ret, rule) } diff --git a/sort/diversity_rule_sort_test.go b/sort/diversity_rule_sort_test.go index f9488db..03cef38 100644 --- a/sort/diversity_rule_sort_test.go +++ b/sort/diversity_rule_sort_test.go @@ -35,10 +35,6 @@ func TestDiversitryRuleSortByIntervalSize(t *testing.T) { }, } - fmt.Println("====sort before====") - for _, item := range items { - t.Log(item) - } context := context.NewRecommendContext() context.Size = 10 sortData := SortData{Data: items, Context: context} @@ -47,9 +43,7 @@ func TestDiversitryRuleSortByIntervalSize(t *testing.T) { result := sortData.Data.([]*module.Item) - fmt.Println("====sort after====") for i, item := range result { - t.Log(item) if i%2 == 0 && item.StringProperty("tag") != "t1" { t.Error("item error") } @@ -109,10 +103,6 @@ func TestDiversitryRuleSortByIntervalSize(t *testing.T) { }, } - fmt.Println("====sort before====") - for _, item := range items { - t.Log(item) - } context := context.NewRecommendContext() context.Size = 10 sortData := SortData{Data: items, Context: context} @@ -121,9 +111,7 @@ func TestDiversitryRuleSortByIntervalSize(t *testing.T) { result := sortData.Data.([]*module.Item) - fmt.Println("====sort after====") for i, item := range result { - t.Log(item) if i%2 == 1 && item.StringProperty("tag") != "t1" { t.Error("item error") } @@ -200,10 +188,6 @@ func TestDiversitryRuleSortByExclusionRule(t *testing.T) { items = append(items, item) } - fmt.Println("====sort before====") - for _, item := range items { - t.Log(item) - } context := context.NewRecommendContext() context.Size = 10 sortData := SortData{Data: items, Context: context} @@ -213,9 +197,7 @@ func TestDiversitryRuleSortByExclusionRule(t *testing.T) { result := sortData.Data.([]*module.Item) assert.Equal(t, 10, len(result)) - fmt.Println("====sort after====") for i, item := range result { - t.Log(item) if i < 5 { assert.Equal(t, "t2", item.StringProperty("tag")) } else { @@ -364,7 +346,7 @@ func TestDiversitryRuleExploreItemSize(t *testing.T) { } } }) - t.Run("diversitry_rule_sort_with_exclusion", func(t *testing.T) { + t.Run("diversitry_rule_sort_with_explore_item_size", func(t *testing.T) { config := recconf.SortConfig{ DiversityRules: []recconf.DiversityRuleConfig{ @@ -444,10 +426,6 @@ func TestDiversitryRuleWeight(t *testing.T) { }, }, } - fmt.Println("====sort before====") - for _, item := range items { - t.Log(item) - } context := context.NewRecommendContext() context.Size = 10 @@ -457,10 +435,8 @@ func TestDiversitryRuleWeight(t *testing.T) { result := sortData.Data.([]*module.Item) - fmt.Println("====sort after====") for i, item := range result { assert.Equal(t, strconv.Itoa(i), string(item.Id)) - t.Log(item) } }) @@ -513,3 +489,260 @@ func TestDiversitryRuleWeight(t *testing.T) { }) } + +func TestDiversitryRuleMulitValue(t *testing.T) { + var items []*module.Item + for i := 0; i < 10; i++ { + item := module.NewItem(strconv.Itoa(i)) + item.Score = float64(i) + item.RetrieveId = "r1" + + item.AddProperty("tag", "A") + items = append(items, item) + } + for i := 10; i < 20; i++ { + item := module.NewItem(strconv.Itoa(i)) + item.Score = float64(i) + item.RetrieveId = "r1" + + item.AddProperty("tag", "A/B") + items = append(items, item) + } + t.Run("diversitry_rule_sort", func(t *testing.T) { + config := recconf.SortConfig{ + + DiversityRules: []recconf.DiversityRuleConfig{ + { + Dimensions: []string{"tag"}, + WindowSize: 5, + //FrequencySize: 1, + IntervalSize: 1, + }, + }, + } + + context := context.NewRecommendContext() + context.Size = 20 + sortData := SortData{Data: items, Context: context} + + NewDiversityRuleSort(config).Sort(&sortData) + + result := sortData.Data.([]*module.Item) + + baseIndex := 0 + for i, item := range result { + if i%2 == 0 { + assert.Equal(t, strconv.Itoa(baseIndex), string(item.Id)) + baseIndex++ + } else { + assert.Equal(t, strconv.Itoa(baseIndex+9), string(item.Id)) + } + } + + }) + t.Run("diversitry_rule_sort_multi_value", func(t *testing.T) { + config := recconf.SortConfig{ + + DiversityRules: []recconf.DiversityRuleConfig{ + { + Dimensions: []string{"tag"}, + WindowSize: 5, + //FrequencySize: 1, + IntervalSize: 1, + }, + }, + MultiValueDimensionConf: []recconf.MultiValueDimensionConfig{ + { + DimensionName: "tag", + Delimiter: "/", + }, + }, + } + + context := context.NewRecommendContext() + context.Size = 20 + sortData := SortData{Data: items, Context: context} + + NewDiversityRuleSort(config).Sort(&sortData) + + result := sortData.Data.([]*module.Item) + + for i, item := range result { + assert.Equal(t, strconv.Itoa(i), string(item.Id)) + } + + }) + t.Run("diversitry_rule_sort_multi_value_v2", func(t *testing.T) { + items = items[:0] + for i := 0; i < 10; i++ { + item := module.NewItem(strconv.Itoa(i)) + item.Score = float64(i) + item.RetrieveId = "r1" + + if i%3 == 0 { + item.AddProperty("tag", "A") + } else if i%3 == 1 { + //item.AddProperty("tag", "A/B") + } else { + item.AddProperty("tag", "B/C") + } + items = append(items, item) + } + config := recconf.SortConfig{ + + DiversityRules: []recconf.DiversityRuleConfig{ + { + Dimensions: []string{"tag"}, + WindowSize: 3, + FrequencySize: 1, + IntervalSize: 1, + }, + }, + MultiValueDimensionConf: []recconf.MultiValueDimensionConfig{ + { + DimensionName: "tag", + Delimiter: "/", + }, + }, + } + + context := context.NewRecommendContext() + context.Size = 10 + sortData := SortData{Data: items, Context: context} + + NewDiversityRuleSort(config).Sort(&sortData) + + result := sortData.Data.([]*module.Item) + + for i, item := range result { + assert.Equal(t, strconv.Itoa(i), string(item.Id)) + } + + }) + t.Run("diversitry_rule_sort_multi_value_v3", func(t *testing.T) { + items = items[:0] + for i := 0; i < 10; i++ { + item := module.NewItem(strconv.Itoa(i)) + item.Score = float64(i) + item.RetrieveId = "r1" + + if i%3 == 0 { + item.AddProperty("tag", "A") + } else if i%3 == 1 { + item.AddProperty("tag", "A/B") + } else { + item.AddProperty("tag", "B/C") + } + if i < 5 { + item.AddProperty("category", strconv.Itoa(0)) + } else { + item.AddProperty("category", strconv.Itoa(1)) + } + items = append(items, item) + } + config := recconf.SortConfig{ + + DiversityRules: []recconf.DiversityRuleConfig{ + { + Dimensions: []string{"category", "tag"}, + WindowSize: 3, + //FrequencySize: 1, + IntervalSize: 1, + }, + }, + MultiValueDimensionConf: []recconf.MultiValueDimensionConfig{ + { + DimensionName: "tag", + Delimiter: "/", + }, + }, + } + fmt.Println("====sort before====") + for _, item := range items { + t.Log(item) + } + + context := context.NewRecommendContext() + context.Size = 10 + sortData := SortData{Data: items, Context: context} + + NewDiversityRuleSort(config).Sort(&sortData) + + result := sortData.Data.([]*module.Item) + + fmt.Println("====sort after====") + for _, item := range result { + //assert.Equal(t, strconv.Itoa(i), string(item.Id)) + t.Log(item) + } + + }) + t.Run("diversitry_rule_sort_multi_value_v4", func(t *testing.T) { + items = items[:0] + for i := 0; i < 10; i++ { + item := module.NewItem(strconv.Itoa(i)) + item.Score = float64(i) + item.RetrieveId = "r1" + + if i%3 == 0 { + item.AddProperty("tag", "A") + } else if i%3 == 1 { + item.AddProperty("tag", "A/B") + } else { + item.AddProperty("tag", "B/C") + } + if i < 3 { + item.AddProperty("category", "c1") + } else if i < 6 { + item.AddProperty("category", "c1#c2") + } else { + item.AddProperty("category", "c3#c4") + } + items = append(items, item) + } + config := recconf.SortConfig{ + + DiversityRules: []recconf.DiversityRuleConfig{ + { + Dimensions: []string{"tag"}, + WindowSize: 3, + //FrequencySize: 1, + IntervalSize: 1, + }, + { + Dimensions: []string{"category"}, + WindowSize: 3, + //FrequencySize: 1, + IntervalSize: 1, + }, + }, + MultiValueDimensionConf: []recconf.MultiValueDimensionConfig{ + { + DimensionName: "tag", + Delimiter: "/", + }, + { + DimensionName: "category", + Delimiter: "#", + }, + }, + } + fmt.Println("====sort before====") + for _, item := range items { + t.Log(item) + } + + context := context.NewRecommendContext() + context.Size = 10 + sortData := SortData{Data: items, Context: context} + + NewDiversityRuleSort(config).Sort(&sortData) + + result := sortData.Data.([]*module.Item) + + assert.Equal(t, string(result[1].Id), strconv.Itoa(8)) + assert.Equal(t, string(result[2].Id), strconv.Itoa(3)) + assert.Equal(t, string(result[3].Id), strconv.Itoa(1)) + + }) +} From fb9a5509995e7e6899ae23587bbcb161ff3edbfe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E9=9D=99?= Date: Thu, 16 Oct 2025 17:04:00 +0800 Subject: [PATCH 2/5] add multi dimension value support for diversity_rule_sort --- sort/diversity_rule.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sort/diversity_rule.go b/sort/diversity_rule.go index d06d6aa..79d967a 100644 --- a/sort/diversity_rule.go +++ b/sort/diversity_rule.go @@ -96,6 +96,8 @@ func (r *DiversityRule) GetWeight() int { } +var _ DiversityRuleInterface = (*DiversityRuleMultiDimension)(nil) + type DiversityRuleMultiDimension struct { DiversityRuleConfig recconf.DiversityRuleConfig DimensionItemMap map[module.ItemId][]any From 7c211b8d2100358bcaad54228c0c454a260cbb4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E9=9D=99?= Date: Thu, 16 Oct 2025 17:45:53 +0800 Subject: [PATCH 3/5] add Negative Cache --- module/realtime_user2item_featurestore_dao.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/module/realtime_user2item_featurestore_dao.go b/module/realtime_user2item_featurestore_dao.go index a857473..8906ae7 100644 --- a/module/realtime_user2item_featurestore_dao.go +++ b/module/realtime_user2item_featurestore_dao.go @@ -203,6 +203,12 @@ func (d *RealtimeUser2ItemFeatureStoreDao) ListItemsByUser(user *User, context * d.cache.Put(triggerId, d.convertItemsToString(items)) } } + //negative cache set, if itemId(triggerId) not have data in feature store, set empty string + for _, itemId := range itemIds { + if _, exist := d.cache.GetIfPresent(itemId); !exist { + d.cache.Put(itemId, "") + } + } } } // sort items From e01d2a60db82a93ec04ac4cd1e6ea99b26d03bb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E9=9D=99?= Date: Fri, 17 Oct 2025 14:08:20 +0800 Subject: [PATCH 4/5] add compress support --- utils/compress/gzip.go | 36 ++++++++++++++++++++++++++++++++++++ utils/compress/lz4.go | 13 +++++++++++++ utils/compress/lz4_test.go | 22 ++++++++++++++++++++++ utils/compress/zstd.go | 17 +++++++++++++++++ web/controller.go | 16 ++++++++++++++++ 5 files changed, 104 insertions(+) create mode 100644 utils/compress/gzip.go create mode 100644 utils/compress/lz4.go create mode 100644 utils/compress/lz4_test.go create mode 100644 utils/compress/zstd.go diff --git a/utils/compress/gzip.go b/utils/compress/gzip.go new file mode 100644 index 0000000..55f18dd --- /dev/null +++ b/utils/compress/gzip.go @@ -0,0 +1,36 @@ +package compress + +import ( + "io" + "sync" + + "github.com/klauspost/compress/gzip" +) + +var gzipReaderPool = sync.Pool{ + // New 函数在池中没有可用对象时被调用，用于创建一个新的对象 + New: func() interface{} { + // 注意：我们直接返回一个 *gzip.Reader 实例。 + // NewReader() 在这里不适用，因为它需要一个 io.Reader 参数。 + // 我们将在使用时通过 Reset() 方法来提供 reader。 + return new(gzip.Reader) + }, +} + +func GzipDecode(reader io.Reader) ([]byte, error) { + gzipReader := gzipReaderPool.Get().(*gzip.Reader) + err := gzipReader.Reset(reader) + if err != nil { + gzipReaderPool.Put(gzipReader) + return nil, err + } + + defer gzipReaderPool.Put(gzipReader) + + decompressed, err := io.ReadAll(gzipReader) + if err != nil { + return nil, err + } + + return decompressed, nil +} diff --git a/utils/compress/lz4.go b/utils/compress/lz4.go new file mode 100644 index 0000000..025326d --- /dev/null +++ b/utils/compress/lz4.go @@ -0,0 +1,13 @@ +package compress + +import ( + "io" + + "github.com/pierrec/lz4" +) + +func LZ4Decode(reader io.Reader) ([]byte, error) { + lz4Reader := lz4.NewReader(reader) + return io.ReadAll(lz4Reader) + +} diff --git a/utils/compress/lz4_test.go b/utils/compress/lz4_test.go new file mode 100644 index 0000000..06f2ff5 --- /dev/null +++ b/utils/compress/lz4_test.go @@ -0,0 +1,22 @@ +package compress + +import ( + "bytes" + "fmt" + "testing" + + "github.com/pierrec/lz4/v4" +) + +func TestLZ4Decode(t *testing.T) { + str := "hello world" + + buf := bytes.NewBuffer(nil) + writer := lz4.NewWriter(buf) + writer.Write([]byte(str)) + writer.Close() + fmt.Println(string(buf.Bytes()), buf.Len()) + + uncompressed, err := LZ4Decode(buf) + fmt.Println(string(uncompressed), err) +} diff --git a/utils/compress/zstd.go b/utils/compress/zstd.go new file mode 100644 index 0000000..421d260 --- /dev/null +++ b/utils/compress/zstd.go @@ -0,0 +1,17 @@ +package compress + +import ( + "io" + + "github.com/klauspost/compress/zstd" +) + +var zstdDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderConcurrency(0)) + +func ZstdDecode(reader io.Reader) ([]byte, error) { + body, err := io.ReadAll(reader) + if err != nil { + return nil, err + } + return zstdDecoder.DecodeAll(body, nil) +} diff --git a/web/controller.go b/web/controller.go index c741808..b531b9d 100644 --- a/web/controller.go +++ b/web/controller.go @@ -8,6 +8,7 @@ import ( "time" "github.com/alibaba/pairec/v2/log" + "github.com/alibaba/pairec/v2/utils/compress" ) const ( @@ -45,6 +46,21 @@ func (c *Controller) cost() int64 { return duration / 1e6 } +func (c *Controller) ReadRequestBody(r *http.Request) ([]byte, error) { + encoding := r.Header.Get("Content-Encoding") + if encoding != "" { + switch encoding { + case "zstd": + return compress.ZstdDecode(r.Body) + case "lz4": + return compress.LZ4Decode(r.Body) + case "gzip": + return compress.GzipDecode(r.Body) + default: + } + } + return io.ReadAll(r.Body) +} func (c *Controller) LogRequestBegin(r *http.Request) { info := fmt.Sprintf("requestId=%s\tevent=begin\turi=%s\taddress=%s\tbody=%s", c.RequestId, r.RequestURI, r.RemoteAddr, string(c.RequestBody)) From 7dbc308d8d474de03782355ef15bcf7087786bb2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E9=9D=99?= Date: Mon, 20 Oct 2025 11:17:16 +0800 Subject: [PATCH 5/5] add response encode support --- utils/compress/gzip.go | 21 +++++++++++++++++---- utils/compress/lz4.go | 15 +++++++++++++++ utils/compress/lz4_test.go | 15 ++++++--------- utils/compress/zstd.go | 15 +++++++++++++++ web/controller.go | 28 ++++++++++++++++++++++++++++ 5 files changed, 81 insertions(+), 13 deletions(-) diff --git a/utils/compress/gzip.go b/utils/compress/gzip.go index 55f18dd..e76b274 100644 --- a/utils/compress/gzip.go +++ b/utils/compress/gzip.go @@ -1,6 +1,7 @@ package compress import ( + "bytes" "io" "sync" @@ -8,11 +9,7 @@ import ( ) var gzipReaderPool = sync.Pool{ - // New 函数在池中没有可用对象时被调用，用于创建一个新的对象 New: func() interface{} { - // 注意：我们直接返回一个 *gzip.Reader 实例。 - // NewReader() 在这里不适用，因为它需要一个 io.Reader 参数。 - // 我们将在使用时通过 Reset() 方法来提供 reader。 return new(gzip.Reader) }, } @@ -34,3 +31,19 @@ func GzipDecode(reader io.Reader) ([]byte, error) { return decompressed, nil } + +func GzipEncode(body []byte) ([]byte, error) { + var buf bytes.Buffer + + gzipWriter := gzip.NewWriter(&buf) + + gzipWriter.Reset(&buf) + if _, err := gzipWriter.Write(body); err != nil { + gzipWriter.Close() + return nil, err + } + if err := gzipWriter.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} diff --git a/utils/compress/lz4.go b/utils/compress/lz4.go index 025326d..900dc84 100644 --- a/utils/compress/lz4.go +++ b/utils/compress/lz4.go @@ -1,6 +1,7 @@ package compress import ( + "bytes" "io" "github.com/pierrec/lz4" @@ -11,3 +12,17 @@ func LZ4Decode(reader io.Reader) ([]byte, error) { return io.ReadAll(lz4Reader) } + +func LZ4Encode(body []byte) ([]byte, error) { + buf := bytes.NewBuffer(nil) + writer := lz4.NewWriter(buf) + if _, err := writer.Write(body); err != nil { + _ = writer.Close() + return nil, err + } + if err := writer.Close(); err != nil { + return nil, err + } + + return buf.Bytes(), nil +} diff --git a/utils/compress/lz4_test.go b/utils/compress/lz4_test.go index 06f2ff5..fe619dc 100644 --- a/utils/compress/lz4_test.go +++ b/utils/compress/lz4_test.go @@ -2,21 +2,18 @@ package compress import ( "bytes" - "fmt" "testing" - "github.com/pierrec/lz4/v4" + "fortio.org/assert" ) func TestLZ4Decode(t *testing.T) { str := "hello world" - buf := bytes.NewBuffer(nil) - writer := lz4.NewWriter(buf) - writer.Write([]byte(str)) - writer.Close() - fmt.Println(string(buf.Bytes()), buf.Len()) + compressData, err := LZ4Encode([]byte(str)) + assert.Equal(t, err, nil) - uncompressed, err := LZ4Decode(buf) - fmt.Println(string(uncompressed), err) + uncompressed, err := LZ4Decode(bytes.NewReader(compressData)) + assert.Equal(t, err, nil) + assert.Equal(t, string(uncompressed), str) } diff --git a/utils/compress/zstd.go b/utils/compress/zstd.go index 421d260..cccc7b8 100644 --- a/utils/compress/zstd.go +++ b/utils/compress/zstd.go @@ -2,6 +2,7 @@ package compress import ( "io" + "sync" "github.com/klauspost/compress/zstd" ) @@ -15,3 +16,17 @@ func ZstdDecode(reader io.Reader) ([]byte, error) { } return zstdDecoder.DecodeAll(body, nil) } + +var zstdWriterPool = sync.Pool{ + New: func() interface{} { + writer, _ := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1)) + return writer + }, +} + +func ZstdEncode(body []byte) ([]byte, error) { + writer := zstdWriterPool.Get().(*zstd.Encoder) + defer zstdWriterPool.Put(writer) + + return writer.EncodeAll(body, make([]byte, 0, len(body))), nil +} diff --git a/web/controller.go b/web/controller.go index b531b9d..9981db8 100644 --- a/web/controller.go +++ b/web/controller.go @@ -93,3 +93,31 @@ func (c *Controller) SendError(w http.ResponseWriter, code int, msg string) { io.WriteString(w, e.ToString()) } + +func (c *Controller) Response(w http.ResponseWriter, r *http.Request, body []byte) { + encoding := r.Header.Get("Accept-Encoding") + if encoding != "" { + switch encoding { + case "zstd": + if buf, err := compress.ZstdEncode(body); err == nil { + w.Header().Add("Content-Encoding", "zstd") + w.Write(buf) + return + } + case "lz4": + if buf, err := compress.LZ4Encode(body); err == nil { + w.Header().Add("Content-Encoding", "lz4") + w.Write(buf) + return + } + case "gzip": + if buf, err := compress.GzipEncode(body); err == nil { + w.Header().Add("Content-Encoding", "gzip") + w.Write(buf) + return + } + default: + } + } + w.Write(body) +}