Move process to create contexts

Signed-off-by: Andrew Thornton <[email protected]>

modules/cron/tasks.go

@@ -80,11 +80,10 @@ func (t *Task) RunWithUser(doer *models.User, config Config) {
 		}
 	}()
 	graceful.GetManager().RunWithShutdownContext(func(baseCtx context.Context) {
-		ctx, cancel := context.WithCancel(baseCtx)
-		defer cancel()
 		pm := process.GetManager()
-		pid := pm.Add(config.FormatMessage(t.Name, "process", doer), cancel)
-		defer pm.Remove(pid)
+		ctx, cancel := pm.AddContext(baseCtx, config.FormatMessage(t.Name, "process", doer))
+		defer cancel()
+
 		if err := t.fun(ctx, doer, config); err != nil {
 			if models.IsErrCancelled(err) {
 				message := err.(models.ErrCancelled).Message

modules/git/blame.go

@@ -25,7 +25,6 @@ type BlamePart struct {
 // BlameReader returns part of file blame one by one
 type BlameReader struct {
 	cmd     *exec.Cmd
-	pid     int64
 	output  io.ReadCloser
 	reader  *bufio.Reader
 	lastSha *string
@@ -100,8 +99,7 @@ func (r *BlameReader) NextPart() (*BlamePart, error) {
 
 // Close BlameReader - don't run NextPart after invoking that
 func (r *BlameReader) Close() error {
-	defer process.GetManager().Remove(r.pid)
-	r.cancel()
+	defer r.cancel()
 
 	_ = r.output.Close()
 
@@ -125,7 +123,8 @@ func CreateBlameReader(ctx context.Context, repoPath, commitID, file string) (*B
 
 func createBlameReader(ctx context.Context, dir string, command ...string) (*BlameReader, error) {
 	// Here we use the provided context - this should be tied to the request performing the blame so that it does not hang around.
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := process.GetManager().AddContext(ctx, fmt.Sprintf("GetBlame [repo_path: %s]", dir))
+
 	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
 	cmd.Dir = dir
 	cmd.Stderr = os.Stderr
@@ -141,13 +140,10 @@ func createBlameReader(ctx context.Context, dir string, command ...string) (*Bla
 		return nil, fmt.Errorf("Start: %v", err)
 	}
 
-	pid := process.GetManager().Add(fmt.Sprintf("GetBlame [repo_path: %s]", dir), cancel)
-
 	reader := bufio.NewReader(stdout)
 
 	return &BlameReader{
 		cmd,
-		pid,
 		stdout,
 		reader,
 		nil,

modules/git/command.go

@@ -143,7 +143,12 @@ func (c *Command) RunWithContext(rc *RunContext) error {
 		log.Debug("%s: %v", rc.Dir, c)
 	}
 
-	ctx, cancel := context.WithTimeout(c.parentContext, rc.Timeout)
+	desc := c.desc
+	if desc == "" {
+		desc = fmt.Sprintf("%s %s [repo_path: %s]", c.name, strings.Join(c.args, " "), rc.Dir)
+	}
+
+	ctx, cancel := process.GetManager().AddContextTimeout(c.parentContext, rc.Timeout, desc)
 	defer cancel()
 
 	cmd := exec.CommandContext(ctx, c.name, c.args...)
@@ -172,13 +177,6 @@ func (c *Command) RunWithContext(rc *RunContext) error {
 		return err
 	}
 
-	desc := c.desc
-	if desc == "" {
-		desc = fmt.Sprintf("%s %s %s [repo_path: %s]", GitExecutable, c.name, strings.Join(c.args, " "), rc.Dir)
-	}
-	pid := process.GetManager().Add(desc, cancel)
-	defer process.GetManager().Remove(pid)
-
 	if rc.PipelineFunc != nil {
 		err := rc.PipelineFunc(ctx, cancel)
 		if err != nil {

modules/git/diff.go

@@ -56,7 +56,7 @@ func GetRepoRawDiffForFile(repo *Repository, startCommit, endCommit string, diff
 		fileArgs = append(fileArgs, "--", file)
 	}
 	// FIXME: graceful: These commands should have a timeout
-	ctx, cancel := context.WithCancel(DefaultContext)
+	ctx, cancel := process.GetManager().AddContext(DefaultContext, fmt.Sprintf("GetRawDiffForFile: [repo_path: %s]", repo.Path))
 	defer cancel()
 
 	var cmd *exec.Cmd
@@ -90,8 +90,6 @@ func GetRepoRawDiffForFile(repo *Repository, startCommit, endCommit string, diff
 	cmd.Dir = repo.Path
 	cmd.Stdout = writer
 	cmd.Stderr = stderr
-	pid := process.GetManager().Add(fmt.Sprintf("GetRawDiffForFile: [repo_path: %s]", repo.Path), cancel)
-	defer process.GetManager().Remove(pid)
 
 	if err = cmd.Run(); err != nil {
 		return fmt.Errorf("Run: %v - %s", err, stderr)

modules/markup/external/external.go

@@ -5,7 +5,6 @@
 package external
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"os"
@@ -107,12 +106,9 @@ func (p *Renderer) Render(ctx *markup.RenderContext, input io.Reader, output io.
 		ctx.Ctx = graceful.GetManager().ShutdownContext()
 	}
 
-	processCtx, cancel := context.WithCancel(ctx.Ctx)
+	processCtx, cancel := process.GetManager().AddContext(ctx.Ctx, fmt.Sprintf("Render [%s] for %s", commands[0], ctx.URLPrefix))
 	defer cancel()
 
-	pid := process.GetManager().Add(fmt.Sprintf("Render [%s] for %s", commands[0], ctx.URLPrefix), cancel)
-	defer process.GetManager().Remove(pid)
-
 	cmd := exec.CommandContext(processCtx, commands[0], args...)
 	cmd.Env = append(
 		os.Environ(),

modules/process/context.go

@@ -0,0 +1,67 @@
+// Copyright 2021 The Gitea Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package process
+
+import (
+	"context"
+)
+
+// Context is a wrapper around context.Context for having the current pid for this context
+type Context struct {
+	context.Context
+	pid int64
+}
+
+// GetPID returns the PID for this context
+func (c *Context) GetPID() int64 {
+	return c.pid
+}
+
+// GetParent returns the parent process context if any
+func (c *Context) GetParent() *Context {
+	return GetContext(c.Context)
+}
+
+func (c *Context) Value(key interface{}) interface{} {
+	if key == ProcessContextKey {
+		return c
+	}
+	return c.Context.Value(key)
+}
+
+// ProcessContextKey is the key under which process contexts are stored
+var ProcessContextKey interface{} = "process-context"
+
+// GetContext will return a process context if one exists
+func GetContext(ctx context.Context) *Context {
+	if pCtx, ok := ctx.(*Context); ok {
+		return pCtx
+	}
+	pCtxInterface := ctx.Value(ProcessContextKey)
+	if pCtxInterface == nil {
+		return nil
+	}
+	if pCtx, ok := pCtxInterface.(*Context); ok {
+		return pCtx
+	}
+	return nil
+}
+
+// GetPID returns the PID for this context
+func GetPID(ctx context.Context) int64 {
+	pCtx := GetContext(ctx)
+	if pCtx == nil {
+		return 0
+	}
+	return pCtx.GetPID()
+}
+
+func GetParentPID(ctx context.Context) int64 {
+	parentPID := int64(0)
+	if parentProcess := GetContext(ctx); parentProcess != nil {
+		parentPID = parentProcess.GetPID()
+	}
+	return parentPID
+}

modules/process/manager.go

@@ -26,11 +26,18 @@ var (
 
 	// DefaultContext is the default context to run processing commands in
 	DefaultContext = context.Background()
+
+	// RecyclePID is the PID number at which we will attempt to recycle PIDs
+	RecyclePID int64 = 1 << 16
+
+	// HuntSize is the size of the hunt for the lowest free PID
+	HuntSize int64 = 512
 )
 
 // Process represents a working process inheriting from Gitea.
 type Process struct {
 	PID         int64 // Process ID, not system one.
+	ParentPID   int64
 	Description string
 	Start       time.Time
 	Cancel      context.CancelFunc
@@ -40,7 +47,9 @@ type Process struct {
 type Manager struct {
 	mutex sync.Mutex
 
-	counter   int64
+	next int64
+	low  int64
+
 	processes map[int64]*Process
 }
 
@@ -49,31 +58,105 @@ func GetManager() *Manager {
 	managerInit.Do(func() {
 		manager = &Manager{
 			processes: make(map[int64]*Process),
+			next:      1,
+			low:       1,
 		}
 	})
 	return manager
 }
 
-// Add a process to the ProcessManager and returns its PID.
-func (pm *Manager) Add(description string, cancel context.CancelFunc) int64 {
+// AddContext create a new context and add it as a process
+func (pm *Manager) AddContext(parent context.Context, description string) (context.Context, context.CancelFunc) {
+	parentPID := GetParentPID(parent)
+
+	ctx, cancel := context.WithCancel(parent)
+
+	pid, cancel := pm.Add(parentPID, description, cancel)
+
+	return &Context{
+		Context: ctx,
+		pid:     pid,
+	}, cancel
+}
+
+// AddContextTimeout create a new context and add it as a process
+func (pm *Manager) AddContextTimeout(parent context.Context, timeout time.Duration, description string) (context.Context, context.CancelFunc) {
+	parentPID := GetParentPID(parent)
+
+	ctx, cancel := context.WithTimeout(parent, timeout)
+
+	pid, cancel := pm.Add(parentPID, description, cancel)
+
+	return &Context{
+		Context: ctx,
+		pid:     pid,
+	}, cancel
+}
+
+// Add create a new process
+func (pm *Manager) Add(parentPID int64, description string, cancel context.CancelFunc) (int64, context.CancelFunc) {
 	pm.mutex.Lock()
-	pid := pm.counter + 1
-	pm.processes[pid] = &Process{
+	pid := pm.nextPID()
+	process := &Process{
 		PID:         pid,
+		ParentPID:   parentPID,
 		Description: description,
 		Start:       time.Now(),
-		Cancel:      cancel,
 	}
-	pm.counter = pid
+
+	process.Cancel = func() {
+		cancel()
+		pm.remove(process)
+	}
+
+	pm.processes[pid] = process
 	pm.mutex.Unlock()
 
-	return pid
+	return pid, process.Cancel
+}
+
+// nextPID will return the next available PID. pm.mutex should already be locked.
+func (pm *Manager) nextPID() int64 {
+	if pm.next > RecyclePID {
+		for i := int64(0); i < HuntSize; i++ {
+			if pm.low >= pm.next {
+				pm.low = 1
+				break
+			}
+			if _, ok := pm.processes[pm.low]; !ok {
+				next := pm.low
+				pm.low++
+				return next
+			}
+			pm.low++
+		}
+	}
+	next := pm.next
+	pm.next++
+	return next
+}
+
+// releasePID will release the PID. pm.mutex should already be locked.
+func (pm *Manager) releasePID(pid int64) {
+	if pid < pm.low+RecyclePID {
+		pm.low = pid
+	}
 }
 
 // Remove a process from the ProcessManager.
 func (pm *Manager) Remove(pid int64) {
 	pm.mutex.Lock()
 	delete(pm.processes, pid)
+	pm.releasePID(pid)
+	pm.mutex.Unlock()
+}
+
+func (pm *Manager) remove(process *Process) {
+	pm.mutex.Lock()
+	if p := pm.processes[process.PID]; p == process {
+		delete(pm.processes, process.PID)
+		pm.releasePID(process.PID)
+	}
 	pm.mutex.Unlock()
 }
 
@@ -134,7 +217,7 @@ func (pm *Manager) ExecDirEnvStdIn(timeout time.Duration, dir, desc string, env
 	stdOut := new(bytes.Buffer)
 	stdErr := new(bytes.Buffer)
 
-	ctx, cancel := context.WithTimeout(DefaultContext, timeout)
+	ctx, cancel := pm.AddContextTimeout(DefaultContext, timeout, desc)
 	defer cancel()
 
 	cmd := exec.CommandContext(ctx, cmdName, args...)
@@ -150,13 +233,11 @@ func (pm *Manager) ExecDirEnvStdIn(timeout time.Duration, dir, desc string, env
 		return "", "", err
 	}
 
-	pid := pm.Add(desc, cancel)
 	err := cmd.Wait()
-	pm.Remove(pid)
 
 	if err != nil {
 		err = &Error{
-			PID:         pid,
+			PID:         GetPID(ctx),
 			Description: desc,
 			Err:         err,
 			CtxErr:      ctx.Err(),