mirror of
https://github.com/openziti/zrok.git
synced 2025-06-28 21:51:26 +02:00
rough proctree implementation (#748)
This commit is contained in:
parent
82d8f4ba2e
commit
42c3ec48b0
59
agent/proctree/impl_posix.go
Normal file
59
agent/proctree/impl_posix.go
Normal file
@ -0,0 +1,59 @@
|
|||||||
|
//go:build !windows
|
||||||
|
|
||||||
|
package proctree
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os/exec"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
func Init(_ string) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func StartChild(tail TailFunction, args ...string) (*Child, error) {
|
||||||
|
cmd := exec.Command(args[0], args[1:]...)
|
||||||
|
|
||||||
|
cld := &Child{
|
||||||
|
TailFunction: tail,
|
||||||
|
cmd: cmd,
|
||||||
|
outStream: make(chan []byte),
|
||||||
|
errStream: make(chan []byte),
|
||||||
|
wg: new(sync.WaitGroup),
|
||||||
|
}
|
||||||
|
|
||||||
|
stdout, err := cmd.StdoutPipe()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
stderr, err := cmd.StderrPipe()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := cmd.Start(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
cld.wg.Add(3)
|
||||||
|
go reader(stdout, cld.outStream, cld.wg)
|
||||||
|
go reader(stderr, cld.errStream, cld.wg)
|
||||||
|
go cld.combiner(cld.wg)
|
||||||
|
|
||||||
|
return cld, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func WaitChild(c *Child) error {
|
||||||
|
c.wg.Wait()
|
||||||
|
if err := c.cmd.Wait(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func StopChild(c *Child) error {
|
||||||
|
if err := c.cmd.Process.Kill(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
79
agent/proctree/impl_windows.go
Executable file
79
agent/proctree/impl_windows.go
Executable file
@ -0,0 +1,79 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
|
package proctree
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/kolesnikovae/go-winjob"
|
||||||
|
"golang.org/x/sys/windows"
|
||||||
|
"os/exec"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
var job *winjob.JobObject
|
||||||
|
|
||||||
|
func Init(name string) error {
|
||||||
|
var err error
|
||||||
|
if job == nil {
|
||||||
|
job, err = winjob.Create(name, winjob.LimitKillOnJobClose, winjob.LimitBreakawayOK)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func StartChild(tail TailFunction, args ...string) (*Child, error) {
|
||||||
|
cmd := exec.Command(args[0], args[1:]...)
|
||||||
|
cmd.SysProcAttr = &windows.SysProcAttr{CreationFlags: windows.CREATE_SUSPENDED}
|
||||||
|
|
||||||
|
cld := &Child{
|
||||||
|
TailFunction: tail,
|
||||||
|
cmd: cmd,
|
||||||
|
outStream: make(chan []byte),
|
||||||
|
errStream: make(chan []byte),
|
||||||
|
wg: new(sync.WaitGroup),
|
||||||
|
}
|
||||||
|
|
||||||
|
stdout, err := cmd.StdoutPipe()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
stderr, err := cmd.StderrPipe()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := cmd.Start(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := job.Assign(cmd.Process); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := winjob.ResumeProcess(cmd.Process.Pid); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
cld.wg.Add(3)
|
||||||
|
go reader(stdout, cld.outStream, cld.wg)
|
||||||
|
go reader(stderr, cld.errStream, cld.wg)
|
||||||
|
go cld.combiner(cld.wg)
|
||||||
|
|
||||||
|
return cld, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func WaitChild(c *Child) error {
|
||||||
|
c.wg.Wait()
|
||||||
|
if err := c.cmd.Wait(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func StopChild(c *Child) error {
|
||||||
|
if err := c.cmd.Process.Kill(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
66
agent/proctree/proctree.go
Executable file
66
agent/proctree/proctree.go
Executable file
@ -0,0 +1,66 @@
|
|||||||
|
package proctree
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"os/exec"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Child struct {
|
||||||
|
TailFunction TailFunction
|
||||||
|
cmd *exec.Cmd
|
||||||
|
outStream chan []byte
|
||||||
|
errStream chan []byte
|
||||||
|
wg *sync.WaitGroup
|
||||||
|
}
|
||||||
|
|
||||||
|
type TailFunction func(data []byte)
|
||||||
|
|
||||||
|
func (c *Child) combiner(wg *sync.WaitGroup) {
|
||||||
|
defer wg.Done()
|
||||||
|
|
||||||
|
outDone := false
|
||||||
|
errDone := false
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case data := <-c.outStream:
|
||||||
|
if data != nil {
|
||||||
|
if c.TailFunction != nil {
|
||||||
|
c.TailFunction(data)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
outDone = true
|
||||||
|
}
|
||||||
|
case data := <-c.errStream:
|
||||||
|
if data != nil {
|
||||||
|
if c.TailFunction != nil {
|
||||||
|
c.TailFunction(data)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
errDone = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if outDone && errDone {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func reader(r io.ReadCloser, o chan []byte, wg *sync.WaitGroup) {
|
||||||
|
defer close(o)
|
||||||
|
defer wg.Done()
|
||||||
|
|
||||||
|
buf := make([]byte, 64*1024)
|
||||||
|
for {
|
||||||
|
n, err := r.Read(buf)
|
||||||
|
if err != nil {
|
||||||
|
if err == io.EOF {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Printf("error reading: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
o <- buf[:n]
|
||||||
|
}
|
||||||
|
}
|
Loading…
x
Reference in New Issue
Block a user