support oom on heap and stack (#79)

Signed-off-by: xiang <xiang13225080@163.com>

cmd/attack/jvm.go

@@ -147,7 +147,7 @@ func NewJVMStressCommand(dep fx.Option, options *core.JVMCommand) *cobra.Command
 	}
 
 	cmd.Flags().IntVarP(&options.CPUCount, "cpu-count", "", 0, "the CPU core number need to use")
-	cmd.Flags().IntVarP(&options.MemorySize, "mem-size", "", 0, "the memory size need to locate, the unit is MB")
+	cmd.Flags().StringVarP(&options.MemoryType, "mem-type", "", "", "the memory type need to locate, the value can be 'stack' or 'heap'")
 
 	return cmd
 }

pkg/core/jvm.go

@@ -61,8 +61,8 @@ type JVMCommand struct {
 	// the CPU core number need to use, only set it when action is stress
 	CPUCount int
 
-	// the memory size need to locate, only set it when action is stress
-	MemorySize int
+	// the memory type need to locate, only set it when action is stress, the value can be 'stack' or 'heap'
+	MemoryType string
 
 	// attach or agent
 	Type string
@@ -80,7 +80,7 @@ type JVMCommand struct {
 
 	StressValueName string
 
-	StressValue int
+	StressValue string
 
 	// btm rule file path
 	RuleFile string
@@ -98,13 +98,20 @@ func (j *JVMCommand) Validate() error {
 	case JVMSubmitType:
 		switch j.Action {
 		case JVMStressAction:
-			if j.CPUCount == 0 && j.MemorySize == 0 {
+			if j.CPUCount == 0 && len(j.MemoryType) == 0 {
 				return errors.New("must set one of cpu-count and mem-size when action is 'stress'")
 			}
 
-			if j.CPUCount > 0 && j.MemorySize > 0 {
+			if j.CPUCount > 0 && len(j.MemoryType) > 0 {
 				return errors.New("inject stress on both CPU and memory is not support now")
 			}
+
+			if len(j.MemoryType) > 0 {
+				if j.MemoryType != "heap" && j.MemoryType != "stack" {
+					return errors.New("memory type should be one of 'heap' and 'stack'")
+				}
+			}
+
 		case JVMGCAction:
 			// do nothing
 		case JVMExceptionAction, JVMReturnAction, JVMLatencyAction:

pkg/core/jvm_test.go

@@ -92,7 +92,7 @@ func TestJVMCommand(t *testing.T) {
 				Type:       JVMSubmitType,
 				Action:     JVMStressAction,
 				CPUCount:   1,
-				MemorySize: 1,
+				MemoryType: "heap",
 			},
 			"inject stress on both CPU and memory is not support now",
 		},

pkg/server/chaosd/jvm.go

@@ -128,11 +128,11 @@ func (j jvmAttack) generateRuleFile(attack *core.JVMCommand) (string, error) {
 			if attack.CPUCount > 0 {
 				attack.StressType = "CPU"
 				attack.StressValueName = "CPUCOUNT"
-				attack.StressValue = attack.CPUCount
+				attack.StressValue = fmt.Sprintf("%d", attack.CPUCount)
 			} else {
 				attack.StressType = "MEMORY"
-				attack.StressValueName = "MEMORYSIZE"
-				attack.StressValue = attack.MemorySize
+				attack.StressValueName = "MEMORYTYPE"
+				attack.StressValue = attack.MemoryType
 			}
 		}
 	}