generate: add timeout options for hooks

Signed-off-by: Ma Shimiao <[email protected]>

Author: Ma Shimiao

cmd/oci-runtime-tool/generate.go

@@ -75,6 +75,12 @@ var generateFlags = []cli.Flag{
 	cli.StringSliceFlag{Name: "mount-bind", Usage: "bind mount directories src:dest[:options...]"},
 	cli.StringFlag{Name: "mount-cgroups", Value: "no", Usage: "mount cgroups (rw,ro,no)"},
 	cli.StringFlag{Name: "output", Usage: "output file (defaults to stdout)"},
+	cli.StringSliceFlag{Name: "poststart", Usage: "set command to run in poststart hooks"},
+	cli.StringSliceFlag{Name: "poststart-timeout", Usage: "set timeout for commands to run in poststart hooks"},
+	cli.StringSliceFlag{Name: "poststop", Usage: "set command to run in poststop hooks"},
+	cli.StringSliceFlag{Name: "poststop-timeout", Usage: "set timeout for commands to run in poststop hooks"},
+	cli.StringSliceFlag{Name: "prestart", Usage: "set command to run in prestart hooks"},
+	cli.StringSliceFlag{Name: "prestart-timeout", Usage: "set timeout for commands to run in prestart hooks"},
 	cli.BoolFlag{Name: "privileged", Usage: "enable privileged container settings"},
 	cli.StringSliceFlag{Name: "process-cap-add", Usage: "add Linux capabilities"},
 	cli.StringSliceFlag{Name: "process-cap-drop", Usage: "drop Linux capabilities"},
@@ -332,14 +338,29 @@ func setupSpec(g *generate.Generator, context *cli.Context) error {
 		}
 	}
 
+<<<<<<< HEAD
 	if context.IsSet("hooks-prestart") {
 		preStartHooks := context.StringSlice("hooks-prestart")
 		for _, hook := range preStartHooks {
 			path, args, err := parseHook(hook)
+=======
+	if context.IsSet("poststart") {
+		postStartHooks := context.StringSlice("poststart")
+		for _, hook := range postStartHooks {
+			path, args := parseHook(hook)
+			g.AddPostStartHook(path, args)
+		}
+	}
+
+	if context.IsSet("poststart-timeout") {
+		postStartTimeouts := context.StringSlice("poststart-timeout")
+		for _, postStartTimeout := range postStartTimeouts {
+			path, timeout, err := parseHookTimeout(postStartTimeout)
+>>>>>>> generate: add timeout options for hooks
 			if err != nil {
-				return err
+				return nil
 			}
-			g.AddPreStartHook(path, args)
+			g.AddPostStartHookTimeout(path, timeout)
 		}
 	}
 
@@ -354,14 +375,40 @@ func setupSpec(g *generate.Generator, context *cli.Context) error {
 		}
 	}
 
+<<<<<<< HEAD
 	if context.IsSet("hooks-poststart") {
 		postStartHooks := context.StringSlice("hooks-poststart")
 		for _, hook := range postStartHooks {
 			path, args, err := parseHook(hook)
+=======
+	if context.IsSet("poststop-timeout") {
+		postStopTimeouts := context.StringSlice("poststop-timeout")
+		for _, postStopTimeout := range postStopTimeouts {
+			path, timeout, err := parseHookTimeout(postStopTimeout)
+>>>>>>> generate: add timeout options for hooks
 			if err != nil {
-				return err
+				return nil
 			}
-			g.AddPostStartHook(path, args)
+			g.AddPostStopHookTimeout(path, timeout)
+		}
+	}
+
+	if context.IsSet("prestart") {
+		preStartHooks := context.StringSlice("prestart")
+		for _, hook := range preStartHooks {
+			path, args := parseHook(hook)
+			g.AddPreStartHook(path, args)
+		}
+	}
+
+	if context.IsSet("prestart-timeout") {
+		preStartTimeouts := context.StringSlice("prestart-timeout")
+		for _, preStartTimeout := range preStartTimeouts {
+			path, timeout, err := parseHookTimeout(preStartTimeout)
+			if err != nil {
+				return nil
+			}
+			g.AddPreStartHookTimeout(path, timeout)
 		}
 	}
 
@@ -641,6 +688,20 @@ func parseHook(s string) (string, []string, error) {
 	return path, args, nil
 }
 
+func parseHookTimeout(s string) (string, int, error) {
+	parts := strings.Split(s, ":")
+	if len(parts) != 2 {
+		return "", 0, fmt.Errorf("invalid format: %s", s)
+	}
+
+	timeout, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return "", 0, err
+	}
+
+	return parts[0], timeout, nil
+}
+
 func parseNetworkPriority(np string) (string, int32, error) {
 	var err error
 

generate/generate.go

@@ -747,6 +747,17 @@ func (g *Generator) AddPreStartHook(path string, args []string) {
 	g.spec.Hooks.Prestart = append(g.spec.Hooks.Prestart, hook)
 }
 
+// AddPreStartHookTimeout adds timeout of a prestart hook into g.spec.Hooks.Prestart.
+func (g *Generator) AddPreStartHookTimeout(path string, timeout int) {
+	g.initSpec()
+	for i, hook := range g.spec.Hooks.Prestart {
+		if hook.Path == path {
+			g.spec.Hooks.Prestart[i].Timeout = &timeout
+			return
+		}
+	}
+}
+
 // ClearPostStopHooks clear g.spec.Hooks.Poststop.
 func (g *Generator) ClearPostStopHooks() {
 	if g.spec == nil {
@@ -765,6 +776,17 @@ func (g *Generator) AddPostStopHook(path string, args []string) {
 	g.spec.Hooks.Poststop = append(g.spec.Hooks.Poststop, hook)
 }
 
+// AddPostStopHookTimeout adds timeout of a poststop hook into g.spec.Hooks.Postsop.
+func (g *Generator) AddPostStopHookTimeout(path string, timeout int) {
+	g.initSpec()
+	for i, hook := range g.spec.Hooks.Poststop {
+		if hook.Path == path {
+			g.spec.Hooks.Poststop[i].Timeout = &timeout
+			return
+		}
+	}
+}
+
 // ClearPostStartHooks clear g.spec.Hooks.Poststart.
 func (g *Generator) ClearPostStartHooks() {
 	if g.spec == nil {
@@ -783,6 +805,17 @@ func (g *Generator) AddPostStartHook(path string, args []string) {
 	g.spec.Hooks.Poststart = append(g.spec.Hooks.Poststart, hook)
 }
 
+// AddPostStartHookTimeout adds timeout of a poststart hook into g.spec.Hooks.Poststart.
+func (g *Generator) AddPostStartHookTimeout(path string, timeout int) {
+	g.initSpec()
+	for i, hook := range g.spec.Hooks.Poststart {
+		if hook.Path == path {
+			g.spec.Hooks.Poststart[i].Timeout = &timeout
+			return
+		}
+	}
+}
+
 // AddTmpfsMount adds a tmpfs mount into g.spec.Mounts.
 func (g *Generator) AddTmpfsMount(dest string, options []string) {
 	mnt := rspec.Mount{