tksadik92/wrenn-releases
1
0
Fork 0
forked from wrenn/wrenn
Code Pull Requests Activity

v0.2.1 (#55)

Browse Source 
Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#55
Co-authored-by: pptx704 <rafeed@omukk.dev>
Co-committed-by: pptx704 <rafeed@omukk.dev>
This commit is contained in:
Rafeed M. Bhuiyan
2026-06-20 22:45:08 +00:00
committed by Rafeed M. Bhuiyan
parent cfc0c52010
commit a08e755e53
53 changed files with 1675 additions and 577 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
internal/envdclient/client.go
View File

@ -25,6 +25,7 @@ type Client struct {
hostIP string
base string
healthURL string
activityURL string
httpClient *http.Client
streamingClient *http.Client
@ -42,6 +43,7 @@ func New(hostIP string) *Client {
hostIP: hostIP,
base: base,
healthURL: base + "/health",
activityURL: base + "/activity",
httpClient: httpClient,
streamingClient: streamingClient,
process: genconnect.NewProcessClient(streamingClient, base),
@ -117,36 +119,17 @@ func (c *Client) Exec(ctx context.Context, cmd string, args []string, opts *Exec
result := &ExecResult{}
for stream.Receive() {
msg := stream.Msg()
if msg.Event == nil {
ev, ok := procEventToStreamEvent(stream.Msg().GetEvent())
if !ok {
continue
}
event := msg.Event.GetEvent()
switch e := event.(type) {
case *envdpb.ProcessEvent_Start:
slog.Debug("process started", "pid", e.Start.GetPid())
case *envdpb.ProcessEvent_Data:
output := e.Data.GetOutput()
switch o := output.(type) {
case *envdpb.ProcessEvent_DataEvent_Stdout:
result.Stdout = append(result.Stdout, o.Stdout...)
case *envdpb.ProcessEvent_DataEvent_Stderr:
result.Stderr = append(result.Stderr, o.Stderr...)
}
case *envdpb.ProcessEvent_End:
result.ExitCode = e.End.GetExitCode()
if e.End.Error != nil {
slog.Debug("process ended with error",
"exit_code", e.End.GetExitCode(),
"error", e.End.GetError(),
)
}
case *envdpb.ProcessEvent_Keepalive:
// Ignore keepalives.
switch ev.Type {
case "stdout":
result.Stdout = append(result.Stdout, ev.Data...)
case "stderr":
result.Stderr = append(result.Stderr, ev.Data...)
case "end":
result.ExitCode = ev.ExitCode
}
}
@ -166,6 +149,76 @@ type ExecStreamEvent struct {
Error string
}
// procEventToStreamEvent converts a raw envd ProcessEvent into an
// ExecStreamEvent. The second return is false for events with no payload to
// forward (nil event, keepalive, unknown data variant) so callers can skip
// them. This is the single decoder shared by Exec, ExecStream and
// ConnectProcess.
func procEventToStreamEvent(pe *envdpb.ProcessEvent) (ExecStreamEvent, bool) {
if pe == nil {
return ExecStreamEvent{}, false
}
switch e := pe.GetEvent().(type) {
case *envdpb.ProcessEvent_Start:
return ExecStreamEvent{Type: "start", PID: e.Start.GetPid()}, true
case *envdpb.ProcessEvent_Data:
switch o := e.Data.GetOutput().(type) {
case *envdpb.ProcessEvent_DataEvent_Stdout:
return ExecStreamEvent{Type: "stdout", Data: o.Stdout}, true
case *envdpb.ProcessEvent_DataEvent_Stderr:
return ExecStreamEvent{Type: "stderr", Data: o.Stderr}, true
}
return ExecStreamEvent{}, false
case *envdpb.ProcessEvent_End:
ev := ExecStreamEvent{Type: "end", ExitCode: e.End.GetExitCode()}
if e.End.Error != nil {
ev.Error = e.End.GetError()
}
return ev, true
}
return ExecStreamEvent{}, false
}
// procEventStream is the subset of a Connect server-stream that pumpProcessEvents
// needs. Both *connect.ServerStreamForClient[StartResponse] and
// [ConnectResponse] satisfy it.
type procEventStream[T any] interface {
Receive() bool
Msg() *T
Err() error
Close() error
}
// pumpProcessEvents drains a process server-stream into ch until the stream ends
// or ctx is cancelled, closing ch on exit. getEvent extracts the ProcessEvent
// from each message so the same loop works for both the Start and Connect RPCs.
func pumpProcessEvents[T any](
ctx context.Context,
stream procEventStream[T],
getEvent func(*T) *envdpb.ProcessEvent,
ch chan<- ExecStreamEvent,
logLabel string,
) {
defer close(ch)
defer stream.Close()
for stream.Receive() {
ev, ok := procEventToStreamEvent(getEvent(stream.Msg()))
if !ok {
continue
}
select {
case ch <- ev:
case <-ctx.Done():
return
}
}
if err := stream.Err(); err != nil && err != io.EOF {
slog.Debug(logLabel, "error", err)
}
}
// ExecStream runs a command inside the sandbox and returns a channel of output events.
// The channel is closed when the process ends or the context is cancelled.
func (c *Client) ExecStream(ctx context.Context, cmd string, args ...string) (<-chan ExecStreamEvent, error) {
@ -184,52 +237,7 @@ func (c *Client) ExecStream(ctx context.Context, cmd string, args ...string) (<-
}
ch := make(chan ExecStreamEvent, 256)
go func() {
defer close(ch)
defer stream.Close()
for stream.Receive() {
msg := stream.Msg()
if msg.Event == nil {
continue
}
var ev ExecStreamEvent
event := msg.Event.GetEvent()
switch e := event.(type) {
case *envdpb.ProcessEvent_Start:
ev = ExecStreamEvent{Type: "start", PID: e.Start.GetPid()}
case *envdpb.ProcessEvent_Data:
output := e.Data.GetOutput()
switch o := output.(type) {
case *envdpb.ProcessEvent_DataEvent_Stdout:
ev = ExecStreamEvent{Type: "stdout", Data: o.Stdout}
case *envdpb.ProcessEvent_DataEvent_Stderr:
ev = ExecStreamEvent{Type: "stderr", Data: o.Stderr}
}
case *envdpb.ProcessEvent_End:
ev = ExecStreamEvent{Type: "end", ExitCode: e.End.GetExitCode()}
if e.End.Error != nil {
ev.Error = e.End.GetError()
}
case *envdpb.ProcessEvent_Keepalive:
continue
}
select {
case ch <- ev:
case <-ctx.Done():
return
}
}
if err := stream.Err(); err != nil && err != io.EOF {
slog.Debug("exec stream error", "error", err)
}
}()
go pumpProcessEvents(ctx, stream, (*envdpb.StartResponse).GetEvent, ch, "exec stream error")
return ch, nil
}
@ -434,7 +442,7 @@ func (c *Client) CancelMemoryPreload(ctx context.Context) error {
// post-restore initialization. sandbox_id and template_id are passed
// so envd can set WRENN_SANDBOX_ID and WRENN_TEMPLATE_ID env vars.
func (c *Client) PostInit(ctx context.Context) error {
return c.PostInitWithDefaults(ctx, "", nil, "", "")
return c.PostInitWithDefaults(ctx, "", nil, "", "", "")
}
// PostInitWithDefaults calls envd's POST /init endpoint with optional default
@ -444,7 +452,7 @@ func (c *Client) PostInit(ctx context.Context) error {
// timestamp and lifecycle_id are always populated: envd uses them to snap
// the guest clock to the host's wall time and to detect post-resume calls
// (which trigger port-forwarder restart + NFS remount).
func (c *Client) PostInitWithDefaults(ctx context.Context, defaultUser string, envVars map[string]string, sandboxID, templateID string) error {
func (c *Client) PostInitWithDefaults(ctx context.Context, defaultUser string, envVars map[string]string, sandboxID, templateID, proxyDomain string) error {
payload := map[string]any{
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"lifecycle_id": uuid.NewString(),
@ -461,6 +469,9 @@ func (c *Client) PostInitWithDefaults(ctx context.Context, defaultUser string, e
if templateID != "" {
payload["template_id"] = templateID
}
if proxyDomain != "" {
payload["proxy_domain"] = proxyDomain
}
var body io.Reader
if len(payload) > 0 {