Skip to content

Commit

Permalink
optimize
Browse files Browse the repository at this point in the history
  • Loading branch information
@rinfx
rinfx committed Dec 5, 2024
1 parent 51f18e7 commit 750621e
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 96 deletions.
29 changes: 19 additions & 10 deletions plugins/wasm-go/examples/custom-log/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,25 +34,34 @@ type CustomLogConfig struct {

// Method 1: write custom log
func writeLog(ctx wrapper.HttpContext) {
ctx.SetCustomLog("question", "当然可以。在Python中，你可以创建一个函数来计算一系列数字的和。下面是一个简单的例子，该函数接受一个数字列表作为输入，并返回它们的总和。\n\n```python\ndef sum_of_numbers(numbers):\n \"\"\"\n 计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = sum(numbers) # 使用Python内置的sum函数计算总和\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers(numbers_list)) # 输出：The sum is: 15\n```\n\n在这段代码中，我们定义了一个名为 `sum_of_numbers` 的函数，它接收一个参数 `numbers`，这是一个包含整数或浮点数的列表。函数内部使用了Python的内置函数 `sum()` 来计算这些数字的总和，并将结果返回。\n\n你也可以手动实现求和逻辑，而不是使用内置的 `sum()` 函数，如下所示：\n\n```python\ndef sum_of_numbers_manual(numbers):\n \"\"\"\n 手动计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = 0\n for number in numbers:\n total_sum += number\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers_manual(numbers_list)) # 输出：The sum is: 15\n```\n\n在这个版本中，我们初始化 `total_sum` 为0，然后遍历列表中的每个元素，并将其加到 `total_sum` 上。最后返回这个累加的结果。这两种方法都可以达到相同的目的，但是使用内置函数通常更简洁且效率更高。")
ctx.SetCustomLog("k2", 2213.22)
ctx.WriteCustomLog()
ctx.SetAttribute("question", "当然可以。在Python中，你可以创建一个函数来计算一系列数字的和。下面是一个简单的例子，该函数接受一个数字列表作为输入，并返回它们的总和。\n\n```python\ndef sum_of_numbers(numbers):\n \"\"\"\n 计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = sum(numbers) # 使用Python内置的sum函数计算总和\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers(numbers_list)) # 输出：The sum is: 15\n```\n\n在这段代码中，我们定义了一个名为 `sum_of_numbers` 的函数，它接收一个参数 `numbers`，这是一个包含整数或浮点数的列表。函数内部使用了Python的内置函数 `sum()` 来计算这些数字的总和，并将结果返回。\n\n你也可以手动实现求和逻辑，而不是使用内置的 `sum()` 函数，如下所示：\n\n```python\ndef sum_of_numbers_manual(numbers):\n \"\"\"\n 手动计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = 0\n for number in numbers:\n total_sum += number\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers_manual(numbers_list)) # 输出：The sum is: 15\n```\n\n在这个版本中，我们初始化 `total_sum` 为0，然后遍历列表中的每个元素，并将其加到 `total_sum` 上。最后返回这个累加的结果。这两种方法都可以达到相同的目的，但是使用内置函数通常更简洁且效率更高。")
ctx.SetAttribute("k2", 2213.22)
ctx.WriteAttributeToLog()
}

// Methods 2: write custom log with specific key
func writeLogWithKey(ctx wrapper.HttpContext, key string) {
ctx.SetCustomLog("k2", 2213.22)
_ = ctx.WriteCustomLogWithKey(key)
ctx.SetCustomLog("k2", 212939.22)
ctx.SetCustomLog("k3", 123)
_ = ctx.WriteCustomLogWithKey(key)
ctx.SetAttribute("k2", 2213.22)
_ = ctx.WriteAttributeToLogWithKey(key)
ctx.SetAttribute("k2", 212939.22)
ctx.SetAttribute("k3", 123)
_ = ctx.WriteAttributeToLogWithKey(key)
}

// Methods 2: write custom log with specific key
func writeTraceAttribute(ctx wrapper.HttpContext) {
ctx.SetAttribute("question", "当然可以。在Python中，你可以创建一个函数来计算一系列数字的和。下面是一个简单的例子，该函数接受一个数字列表作为输入，并返回它们的总和。\n\n```python\ndef sum_of_numbers(numbers):\n \"\"\"\n 计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = sum(numbers) # 使用Python内置的sum函数计算总和\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers(numbers_list)) # 输出：The sum is: 15\n```\n\n在这段代码中，我们定义了一个名为 `sum_of_numbers` 的函数，它接收一个参数 `numbers`，这是一个包含整数或浮点数的列表。函数内部使用了Python的内置函数 `sum()` 来计算这些数字的总和，并将结果返回。\n\n你也可以手动实现求和逻辑，而不是使用内置的 `sum()` 函数，如下所示：\n\n```python\ndef sum_of_numbers_manual(numbers):\n \"\"\"\n 手动计算列表中所有数字的和。\n \n 参数:\n numbers (list of int or float): 一个包含数字的列表。\n \n 返回:\n int or float: 列表中所有数字的总和。\n \"\"\"\n total_sum = 0\n for number in numbers:\n total_sum += number\n return total_sum\n\n# 示例使用\nnumbers_list = [1, 2, 3, 4, 5]\nprint(\"The sum is:\", sum_of_numbers_manual(numbers_list)) # 输出：The sum is: 15\n```\n\n在这个版本中，我们初始化 `total_sum` 为0，然后遍历列表中的每个元素，并将其加到 `total_sum` 上。最后返回这个累加的结果。这两种方法都可以达到相同的目的，但是使用内置函数通常更简洁且效率更高。")
ctx.SetAttribute("k2", 2213.22)
ctx.WriteAttributeToTrace()
}

func onHttpRequestHeaders(ctx wrapper.HttpContext, config CustomLogConfig, log wrapper.Log) types.Action {
if rand.Intn(10)%2 == 1 {
if rand.Intn(10)%3 == 1 {
writeLog(ctx)
} else {
} else if rand.Intn(10)%3 == 2 {
writeLogWithKey(ctx, "ai_log")
} else {
writeTraceAttribute(ctx)
}
return types.ActionContinue
}
86 changes: 69 additions & 17 deletions plugins/wasm-go/pkg/wrapper/plugin_wrapper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
package wrapper

import (
"encoding/json"
"fmt"
"strconv"
"time"
"unsafe"
Expand All @@ -32,6 +34,12 @@ func sched_yield() int32 {
return 0
}

const (
CustomLogKey = "custom_log"
AILogKey = "ai_log"
TraceSpanTagPrefix = "trace_span_tag."
)

type HttpContext interface {
Scheme() string
Host() string
Expand All @@ -41,12 +49,14 @@ type HttpContext interface {
GetContext(key string) interface{}
GetBoolContext(key string, defaultValue bool) bool
GetStringContext(key, defaultValue string) string
GetCustomLog(key string) interface{}
SetCustomLog(key string, value interface{})
GetAttribute(key string) interface{}
SetAttribute(key string, value interface{})
// You can call this function to set custom log
WriteCustomLog() error
WriteAttributeToLog() error
// You can call this function to set custom log with your specific key
WriteCustomLogWithKey(key string) error
WriteAttributeToLogWithKey(key string) error
// You can call this function to set custom trace span attribute
WriteAttributeToTrace() error
// If the onHttpRequestBody handle is not set, the request body will not be read by default
DontReadRequestBody()
// If the onHttpResponseBody handle is not set, the request body will not be read by default
Expand Down Expand Up @@ -269,10 +279,10 @@ func (ctx *CommonPluginCtx[PluginConfig]) OnTick() {

func (ctx *CommonPluginCtx[PluginConfig]) NewHttpContext(contextID uint32) types.HttpContext {
httpCtx := &CommonHttpCtx[PluginConfig]{
plugin: ctx,
contextID: contextID,
userContext: map[string]interface{}{},
customLog: map[string]interface{}{},
plugin: ctx,
contextID: contextID,
userContext: map[string]interface{}{},
userAttribute: map[string]interface{}{},
}
if ctx.vm.onHttpRequestBody != nil || ctx.vm.onHttpStreamingRequestBody != nil {
httpCtx.needRequestBody = true
Expand Down Expand Up @@ -302,7 +312,7 @@ type CommonHttpCtx[PluginConfig any] struct {
responseBodySize int
contextID uint32
userContext map[string]interface{}
customLog map[string]interface{}
userAttribute map[string]interface{}
}

func (ctx *CommonHttpCtx[PluginConfig]) SetContext(key string, value interface{}) {
Expand All @@ -313,20 +323,62 @@ func (ctx *CommonHttpCtx[PluginConfig]) GetContext(key string) interface{} {
return ctx.userContext[key]
}

func (ctx *CommonHttpCtx[PluginConfig]) SetCustomLog(key string, value interface{}) {
ctx.customLog[key] = value
func (ctx *CommonHttpCtx[PluginConfig]) SetAttribute(key string, value interface{}) {
ctx.userAttribute[key] = value
}

func (ctx *CommonHttpCtx[PluginConfig]) GetCustomLog(key string) interface{} {
return ctx.customLog[key]
func (ctx *CommonHttpCtx[PluginConfig]) GetAttribute(key string) interface{} {
return ctx.userAttribute[key]
}

func (ctx *CommonHttpCtx[PluginConfig]) WriteCustomLog() error {
return extendAccessLog(CustomLogKey, ctx.customLog)
func (ctx *CommonHttpCtx[PluginConfig]) WriteAttributeToLog() error {
return ctx.WriteAttributeToLogWithKey(CustomLogKey)
}

func (ctx *CommonHttpCtx[PluginConfig]) WriteCustomLogWithKey(key string) error {
return extendAccessLog(key, ctx.customLog)
func (ctx *CommonHttpCtx[PluginConfig]) WriteAttributeToLogWithKey(key string) error {
// e.g. {\"field1\":\"value1\",\"field2\":\"value2\"}
preMarshalledJsonLogStr, _ := proxywasm.GetProperty([]string{key})
newAttributeMap := map[string]interface{}{}
if string(preMarshalledJsonLogStr) != "" {
// e.g. {"field1":"value1","field2":"value2"}
preJsonLogStr := unmarshalStr(fmt.Sprintf(`"%s"`, string(preMarshalledJsonLogStr)))
err := json.Unmarshal([]byte(preJsonLogStr), &newAttributeMap)
if err != nil {
ctx.plugin.vm.log.Warnf("Unmarshal failed, will overwrite %s, pre value is: %s", key, string(preMarshalledJsonLogStr))
return err
}
}
// update customLog
for k, v := range ctx.userAttribute {
newAttributeMap[k] = v
}
// e.g. {"field1":"value1","field2":2,"field3":"value3"}
jsonStr, _ := json.Marshal(newAttributeMap)
// e.g. {\"field1\":\"value1\",\"field2\":2,\"field3\":\"value3\"}
marshalledJsonStr := marshalStr(string(jsonStr))
if err := proxywasm.SetProperty([]string{key}, []byte(marshalledJsonStr)); err != nil {
ctx.plugin.vm.log.Warnf("failed to set %s in filter state, raw is %s, err is %v", key, marshalledJsonStr, err)
return err
}
return nil
}

func (ctx *CommonHttpCtx[PluginConfig]) WriteAttributeToTrace() error {
for k, v := range ctx.userAttribute {
traceSpanTag := TraceSpanTagPrefix + k
traceSpanValue := fmt.Sprint(v)
var err error
if traceSpanValue != "" {
err = proxywasm.SetProperty([]string{traceSpanTag}, []byte(traceSpanValue))
} else {
err = fmt.Errorf("value of %s is empty", traceSpanTag)
}
if err != nil {
ctx.plugin.vm.log.Warnf("Failed to set trace attribute - %s: %s", traceSpanTag, traceSpanValue)
return err
}
}
return nil
}

func (ctx *CommonHttpCtx[PluginConfig]) GetBoolContext(key string, defaultValue bool) bool {
Expand Down
69 changes: 0 additions & 69 deletions plugins/wasm-go/pkg/wrapper/property_wrapper.go
View file

This file was deleted.

36 changes: 36 additions & 0 deletions plugins/wasm-go/pkg/wrapper/utils.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
package wrapper

import (
"encoding/json"

"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
"github.com/tidwall/gjson"
)

func unmarshalStr(marshalledJsonStr string) string {
// e.g. "{\"field1\":\"value1\",\"field2\":\"value2\"}"
var jsonStr string
err := json.Unmarshal([]byte(marshalledJsonStr), &jsonStr)
if err != nil {
proxywasm.LogErrorf("failed to unmarshal json string, raw string is: %s, err is: %v", marshalledJsonStr, err)
return ""
}
// e.g. {"field1":"value1","field2":"value2"}
return jsonStr
}

func marshalStr(raw string) string {
// e.g. {"field1":"value1","field2":"value2"}
helper := map[string]string{
"placeholder": raw,
}
marshalledHelper, _ := json.Marshal(helper)
marshalledRaw := gjson.GetBytes(marshalledHelper, "placeholder").Raw
if len(marshalledRaw) >= 2 {
// e.g. {\"field1\":\"value1\",\"field2\":\"value2\"}
return marshalledRaw[1 : len(marshalledRaw)-1]
} else {
proxywasm.LogErrorf("failed to marshal json string, raw string is: %s", raw)
return ""
}
}

0 comments on commit 750621e

Please sign in to comment.