Skip to content

chore: add backed reader, writer and pipe #19225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: mike/immortal-streams-backed-base
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ import (
"github.com/coder/coder/v2/agent/agentexec"
"github.com/coder/coder/v2/agent/agentscripts"
"github.com/coder/coder/v2/agent/agentssh"
"github.com/coder/coder/v2/agent/immortalstreams"
"github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/agent/proto/resourcesmonitor"
"github.com/coder/coder/v2/agent/reconnectingpty"
Expand Down Expand Up @@ -280,6 +281,9 @@ type agent struct {
devcontainers bool
containerAPIOptions []agentcontainers.Option
containerAPI *agentcontainers.API

// Immortal streams
immortalStreamsManager *immortalstreams.Manager
}

func (a *agent) TailnetConn() *tailnet.Conn {
Expand Down Expand Up @@ -347,6 +351,9 @@ func (a *agent) init() {

a.containerAPI = agentcontainers.NewAPI(a.logger.Named("containers"), containerAPIOpts...)

// Initialize immortal streams manager
a.immortalStreamsManager = immortalstreams.New(a.logger.Named("immortal-streams"), &net.Dialer{})

a.reconnectingPTYServer = reconnectingpty.NewServer(
a.logger.Named("reconnecting-pty"),
a.sshServer,
Expand Down Expand Up @@ -1930,6 +1937,12 @@ func (a *agent) Close() error {
a.logger.Error(a.hardCtx, "container API close", slog.Error(err))
}

if a.immortalStreamsManager != nil {
if err := a.immortalStreamsManager.Close(); err != nil {
a.logger.Error(a.hardCtx, "immortal streams manager close", slog.Error(err))
}
}

// Wait for the graceful shutdown to complete, but don't wait forever so
// that we don't break user expectations.
go func() {
Expand Down
7 changes: 7 additions & 0 deletions agent/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"github.com/go-chi/chi/v5"
"github.com/google/uuid"

"github.com/coder/coder/v2/coderd/agentapi"
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/codersdk"
)
Expand Down Expand Up @@ -66,6 +67,12 @@ func (a *agent) apiHandler() http.Handler {
r.Get("/debug/manifest", a.HandleHTTPDebugManifest)
r.Get("/debug/prometheus", promHandler.ServeHTTP)

// Mount immortal streams API
if a.immortalStreamsManager != nil {
immortalStreamsHandler := agentapi.NewImmortalStreamsHandler(a.logger, a.immortalStreamsManager)
r.Mount("/api/v0/immortal-stream", immortalStreamsHandler.Routes())
}

return r
}

Expand Down
235 changes: 235 additions & 0 deletions agent/immortalstreams/manager.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,235 @@
package immortalstreams

import (
"context"
"fmt"
"io"
"net"
"sync"
"time"

"github.com/google/uuid"
"github.com/moby/moby/pkg/namesgenerator"
"golang.org/x/xerrors"

"cdr.dev/slog"
"github.com/coder/coder/v2/codersdk"
)

const (
// MaxStreams is the maximum number of immortal streams allowed per agent
MaxStreams = 32
// BufferSize is the size of the ring buffer for each stream (64 MiB)
BufferSize = 64 * 1024 * 1024
)

// Manager manages immortal streams for an agent
type Manager struct {
logger slog.Logger

mu sync.RWMutex
streams map[uuid.UUID]*Stream

// dialer is used to dial local services
dialer Dialer
}

// Dialer dials a local service
type Dialer interface {
DialContext(ctx context.Context, network, address string) (net.Conn, error)
}

// New creates a new immortal streams manager
func New(logger slog.Logger, dialer Dialer) *Manager {
return &Manager{
logger: logger,
streams: make(map[uuid.UUID]*Stream),
dialer: dialer,
}
}

// CreateStream creates a new immortal stream
func (m *Manager) CreateStream(ctx context.Context, port int) (*codersdk.ImmortalStream, error) {
m.mu.Lock()
defer m.mu.Unlock()

// Check if we're at the limit
if len(m.streams) >= MaxStreams {
// Try to evict a disconnected stream
evicted := m.evictOldestDisconnectedLocked()
if !evicted {
return nil, xerrors.New("too many immortal streams")
}
}

// Dial the local service
addr := fmt.Sprintf("localhost:%d", port)
conn, err := m.dialer.DialContext(ctx, "tcp", addr)
if err != nil {
if isConnectionRefused(err) {
return nil, xerrors.Errorf("the connection was refused")
}
return nil, xerrors.Errorf("dial local service: %w", err)
}

// Create the stream
id := uuid.New()
name := namesgenerator.GetRandomName(0)
stream := NewStream(
id,
name,
port,
m.logger.With(slog.F("stream_id", id), slog.F("stream_name", name)),
BufferSize,
)

// Start the stream
if err := stream.Start(conn); err != nil {
_ = conn.Close()
return nil, xerrors.Errorf("start stream: %w", err)
}

m.streams[id] = stream

return &codersdk.ImmortalStream{
ID: id,
Name: name,
TCPPort: port,
CreatedAt: stream.createdAt,
LastConnectionAt: stream.createdAt,
}, nil
}

// GetStream returns a stream by ID
func (m *Manager) GetStream(id uuid.UUID) (*Stream, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
stream, ok := m.streams[id]
return stream, ok
}

// ListStreams returns all streams
func (m *Manager) ListStreams() []codersdk.ImmortalStream {
m.mu.RLock()
defer m.mu.RUnlock()

streams := make([]codersdk.ImmortalStream, 0, len(m.streams))
for _, stream := range m.streams {
streams = append(streams, stream.ToAPI())
}
return streams
}

// DeleteStream deletes a stream by ID
func (m *Manager) DeleteStream(id uuid.UUID) error {
m.mu.Lock()
defer m.mu.Unlock()

stream, ok := m.streams[id]
if !ok {
return xerrors.New("stream not found")
}

if err := stream.Close(); err != nil {
m.logger.Warn(context.Background(), "failed to close stream", slog.Error(err))
}

delete(m.streams, id)
return nil
}

// Close closes all streams
func (m *Manager) Close() error {
m.mu.Lock()
defer m.mu.Unlock()

var firstErr error
for id, stream := range m.streams {
if err := stream.Close(); err != nil && firstErr == nil {
firstErr = err
}
delete(m.streams, id)
}
return firstErr
}

// evictOldestDisconnectedLocked evicts the oldest disconnected stream
// Must be called with mu held
func (m *Manager) evictOldestDisconnectedLocked() bool {
var (
oldestID uuid.UUID
oldestDisconnected time.Time
found bool
)

for id, stream := range m.streams {
if stream.IsConnected() {
continue
}

disconnectedAt := stream.LastDisconnectionAt()

// Prioritize streams that have actually been disconnected over never-connected streams
switch {
case !found:
oldestID = id
oldestDisconnected = disconnectedAt
found = true
case disconnectedAt.IsZero() && !oldestDisconnected.IsZero():
// Keep the current choice (it was actually disconnected)
continue
case !disconnectedAt.IsZero() && oldestDisconnected.IsZero():
// Prefer this stream (it was actually disconnected) over never-connected
oldestID = id
oldestDisconnected = disconnectedAt
case !disconnectedAt.IsZero() && !oldestDisconnected.IsZero():
// Both were actually disconnected, pick the oldest
if disconnectedAt.Before(oldestDisconnected) {
oldestID = id
oldestDisconnected = disconnectedAt
}
}
// If both are zero time, keep the first one found
}

if !found {
return false
}

// Close and remove the oldest disconnected stream
if stream, ok := m.streams[oldestID]; ok {
m.logger.Info(context.Background(), "evicting oldest disconnected stream",
slog.F("stream_id", oldestID),
slog.F("stream_name", stream.name),
slog.F("disconnected_at", oldestDisconnected))

if err := stream.Close(); err != nil {
m.logger.Warn(context.Background(), "failed to close evicted stream", slog.Error(err))
}
delete(m.streams, oldestID)
}

return true
}

// HandleConnection handles a new connection for an existing stream
func (m *Manager) HandleConnection(id uuid.UUID, conn io.ReadWriteCloser, readSeqNum uint64) error {
m.mu.RLock()
stream, ok := m.streams[id]
m.mu.RUnlock()

if !ok {
return xerrors.New("stream not found")
}

return stream.HandleReconnect(conn, readSeqNum)
}

// isConnectionRefused checks if an error is a connection refused error
func isConnectionRefused(err error) bool {
var opErr *net.OpError
if xerrors.As(err, &opErr) {
return opErr.Op == "dial"
}
return false
}
Loading
Loading
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy