From 379097c51647bdc73e4231396fbc7181fc15fb6f Mon Sep 17 00:00:00 2001
From: xucong053 <xucong053@outlook.com>
Date: Thu, 26 May 2022 10:51:59 +0800
Subject: [PATCH] suport for distributing tasks that contain plugins

---
 hrp/boomer.go                      | 27 +++++++++++++++++++++++++++
 hrp/config.go                      |  9 ++++++++-
 hrp/internal/boomer/client_grpc.go |  4 ++--
 hrp/runner.go                      | 14 ++++++++++++++
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/hrp/boomer.go b/hrp/boomer.go
index 8a076a4a..87c6c809 100644
--- a/hrp/boomer.go
+++ b/hrp/boomer.go
@@ -1,7 +1,11 @@
 package hrp
 
 import (
+	"fmt"
+	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"sync"
 	"time"
 
@@ -194,7 +198,30 @@ func (b *HRPBoomer) runTasks(testCases []*TCase, profile *boomer.Profile) {
 		tesecase, err := tc.toTestCase()
 		if err != nil {
 			log.Error().Err(err).Msg("failed to load testcases")
+			return
 		}
+		// create temp dir to save testcase
+		tempDir, err := ioutil.TempDir("", "hrp_testcases")
+		if err != nil {
+			log.Error().Err(err).Msg("failed to save testcases")
+			return
+		}
+
+		tesecase.Config.Path = filepath.Join(tempDir, "test-case.json")
+		if tesecase.Config.PluginSetting != nil {
+			tesecase.Config.PluginSetting.Path = filepath.Join(tempDir, fmt.Sprintf("debugtalk.%s", tesecase.Config.PluginSetting.Type))
+			err = builtin.Bytes2File(tesecase.Config.PluginSetting.Content, tesecase.Config.PluginSetting.Path)
+			if err != nil {
+				log.Error().Err(err).Msg("failed to save plugin file")
+				return
+			}
+		}
+		err = builtin.Dump2JSON(tesecase, tesecase.Config.Path)
+		if err != nil {
+			log.Error().Err(err).Msg("failed to dump testcases")
+			return
+		}
+
 		testcases = append(testcases, tesecase)
 	}
 	b.SetProfile(profile)
diff --git a/hrp/config.go b/hrp/config.go
index 06930564..3ee1264a 100644
--- a/hrp/config.go
+++ b/hrp/config.go
@@ -32,7 +32,8 @@ type TConfig struct {
 	Timeout           float64                `json:"timeout,omitempty" yaml:"timeout,omitempty"` // global timeout in seconds
 	Export            []string               `json:"export,omitempty" yaml:"export,omitempty"`
 	Weight            int                    `json:"weight,omitempty" yaml:"weight,omitempty"`
-	Path              string                 `json:"path,omitempty" yaml:"path,omitempty"` // testcase file path
+	Path              string                 `json:"path,omitempty" yaml:"path,omitempty"`     // testcase file path
+	PluginSetting     *PluginConfig          `json:"plugin,omitempty" yaml:"plugin,omitempty"` // plugin config
 }
 
 // WithVariables sets variables for current testcase.
@@ -172,3 +173,9 @@ const (
 )
 
 var thinkTimeDefaultRandom = map[string]float64{"min_percentage": 0.5, "max_percentage": 1.5}
+
+type PluginConfig struct {
+	Path    string
+	Type    string // bin、so、py
+	Content []byte
+}
diff --git a/hrp/internal/boomer/client_grpc.go b/hrp/internal/boomer/client_grpc.go
index 8082074d..892f9c89 100644
--- a/hrp/internal/boomer/client_grpc.go
+++ b/hrp/internal/boomer/client_grpc.go
@@ -74,12 +74,12 @@ func newClient(masterHost string, masterPort int, identity string) (client *grpc
 
 func (c *grpcClient) connect() (err error) {
 	addr := fmt.Sprintf("%v:%v", c.masterHost, c.masterPort)
-	c.config.conn, err = grpc.Dial(addr, grpc.WithInsecure())
+	c.config.conn, err = grpc.Dial(addr, grpc.WithInsecure(), grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(1024*1024*1024)))
 	if err != nil {
 		log.Error().Err(err).Msg("failed to connect")
 		return err
 	}
-
+	grpc.MaxCallRecvMsgSize(32 * 10e9)
 	go c.recv()
 	go c.send()
 
diff --git a/hrp/runner.go b/hrp/runner.go
index 12272d4b..63a3098a 100644
--- a/hrp/runner.go
+++ b/hrp/runner.go
@@ -7,6 +7,7 @@ import (
 	"net/http/cookiejar"
 	"net/url"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -16,6 +17,7 @@ import (
 	"github.com/rs/zerolog/log"
 	"golang.org/x/net/http2"
 
+	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
 	"github.com/httprunner/httprunner/v4/hrp/internal/sdk"
 )
 
@@ -279,6 +281,18 @@ func (r *HRPRunner) newCaseRunner(testcase *TestCase) (*testCaseRunner, error) {
 		timeout := time.Duration(runner.testCase.Config.Timeout*1000) * time.Millisecond
 		runner.hrpRunner.SetTimeout(timeout)
 	}
+	if plugin.Path() != "" {
+		pluginContent, err := builtin.ReadFile(plugin.Path())
+		if err != nil {
+			return nil, err
+		}
+		tp := strings.Split(plugin.Path(), ".")
+		runner.parsedConfig.PluginSetting = &PluginConfig{
+			Path:    plugin.Path(),
+			Content: pluginContent,
+			Type:    tp[len(tp)-1],
+		}
+	}
 
 	return runner, nil
 }