optimize: high cpu

common/utils.go

@@ -13,6 +13,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	"unsafe"
 )
 
 func OpenBrowser(url string) {
@@ -142,7 +143,7 @@ func init() {
 }
 
 func GenerateKey() string {
-	rand.Seed(time.Now().UnixNano())
+	//rand.Seed(time.Now().UnixNano())
 	key := make([]byte, 48)
 	for i := 0; i < 16; i++ {
 		key[i] = keyChars[rand.Intn(len(keyChars))]
@@ -159,7 +160,7 @@ func GenerateKey() string {
 }
 
 func GetRandomString(length int) string {
-	rand.Seed(time.Now().UnixNano())
+	//rand.Seed(time.Now().UnixNano())
 	key := make([]byte, length)
 	for i := 0; i < length; i++ {
 		key[i] = keyChars[rand.Intn(len(keyChars))]
@@ -216,3 +217,10 @@ func StringsContains(strs []string, str string) bool {
 	}
 	return false
 }
+
+// []byte only read, panic on append
+func StringToByteSlice(s string) []byte {
+	tmp1 := (*[2]uintptr)(unsafe.Pointer(&s))
+	tmp2 := [3]uintptr{tmp1[0], tmp1[1], tmp1[1]}
+	return *(*[]byte)(unsafe.Pointer(&tmp2))
+}

controller/relay-openai.go

@@ -12,7 +12,7 @@ import (
 )
 
 func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) {
-	responseText := ""
+	var responseTextBuilder strings.Builder
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
 		if atEOF && len(data) == 0 {
@@ -29,6 +29,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	dataChan := make(chan string)
 	stopChan := make(chan bool)
 	go func() {
+		var streamItems []string
 		for scanner.Scan() {
 			data := scanner.Text()
 			if len(data) < 6 { // ignore blank line or wrong format
@@ -40,27 +41,33 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 			dataChan <- data
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
+				streamItems = append(streamItems, data)
+			}
+		}
+		streamResp := "[" + strings.Join(streamItems, ",") + "]"
 		switch relayMode {
 		case RelayModeChatCompletions:
-					var streamResponse ChatCompletionsStreamResponse
+			var streamResponses []ChatCompletionsStreamResponseSimple
-					err := json.Unmarshal([]byte(data), &streamResponse)
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
 			if err != nil {
 				common.SysError("error unmarshalling stream response: " + err.Error())
-						continue // just ignore the error
+				return // just ignore the error
 			}
+			for _, streamResponse := range streamResponses {
 				for _, choice := range streamResponse.Choices {
-						responseText += choice.Delta.Content
+					responseTextBuilder.WriteString(choice.Delta.Content)
+				}
 			}
 		case RelayModeCompletions:
-					var streamResponse CompletionsStreamResponse
+			var streamResponses []CompletionsStreamResponse
-					err := json.Unmarshal([]byte(data), &streamResponse)
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
 			if err != nil {
 				common.SysError("error unmarshalling stream response: " + err.Error())
-						continue
+				return // just ignore the error
 			}
+			for _, streamResponse := range streamResponses {
 				for _, choice := range streamResponse.Choices {
-						responseText += choice.Text
+					responseTextBuilder.WriteString(choice.Text)
-					}
 				}
 			}
 		}
@@ -85,7 +92,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	if err != nil {
 		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
 	}
-	return nil, responseText
+	return nil, responseTextBuilder.String()
 }
 
 func openaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*OpenAIErrorWithStatusCode, *Usage) {

controller/relay.go

@@ -189,6 +189,10 @@ type ChatCompletionsStreamResponse struct {
 	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
 }
 
+type ChatCompletionsStreamResponseSimple struct {
+	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
+}
+
 type CompletionsStreamResponse struct {
 	Choices []struct {
 		Text         string `json:"text"`

main.go

@@ -6,6 +6,8 @@ import (
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-contrib/sessions/cookie"
 	"github.com/gin-gonic/gin"
+	"log"
+	"net/http"
 	"one-api/common"
 	"one-api/controller"
 	"one-api/middleware"
@@ -13,6 +15,8 @@ import (
 	"one-api/router"
 	"os"
 	"strconv"
+
+	_ "net/http/pprof"
 )
 
 //go:embed web/build
@@ -85,6 +89,9 @@ func main() {
 	}
 
 	if os.Getenv("ENABLE_PPROF") == "true" {
+		go func() {
+			log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
+		}()
 		go common.Monitor()
 		common.SysLog("pprof enabled")
 	}