package handlers import ( "sync" "time" "github.com/mimecast/dtail/internal/io/dlog" user "github.com/mimecast/dtail/internal/user/server" ) const ( defaultOutputChannelBufferSize = 1000 defaultOutputFlushTimeout = 2 * time.Second defaultOutputFlushPollInterval = 10 * time.Millisecond defaultOutputReadRetryInterval = time.Millisecond defaultOutputEOFAckQuietPeriod = 50 * time.Millisecond defaultOutputEOFAckTimeout = 2 * time.Second ) type outputManagerConfig struct { channelBufferSize int flushTimeout time.Duration flushPollInterval time.Duration readRetryInterval time.Duration eofAckQuietPeriod time.Duration } // outputManager coordinates output-mode state between command goroutines // (enable/signalEOF/flush/waitForEOFAck, spawned per command by the server // handler) and the session output goroutine (io.Copy -> baseHandler.Read -> // tryRead). These run concurrently, so all mutable state is guarded by mu: // mode, lines, buffer, eof, eofAck, eofEmptySince and epoch must only be // accessed while holding mu. Channel *operations* (send/receive/close/len) // are safe on a snapshot taken under mu; only the field reads/writes need // the lock. // // The configuration fields (channelBufferSize, flushTimeout, ...) are // deliberately not guarded: configure() runs exactly once from the handler // constructor before any goroutine can touch the manager, so goroutine // creation establishes the necessary happens-before edge. type outputManager struct { mu sync.Mutex mode bool lines chan []byte buffer []byte eof chan struct{} eofAck chan struct{} // epoch is bumped by every enable() call and identifies the newest batch // that joined the output session. signalEOF only closes the EOF channel // when the signaler's captured epoch is still current, so a command that // decided "the batch is over" before another command joined cannot EOF // the newcomer's output (see signalEOF for the full protocol). epoch uint64 channelBufferSize int flushTimeout time.Duration flushPollInterval time.Duration readRetryInterval time.Duration eofAckQuietPeriod time.Duration eofEmptySince time.Time } // configure sets the tunables. It must be called before the manager is used // concurrently (i.e. from the handler constructor); see the struct comment // for why the config fields need no locking. func (t *outputManager) configure(cfg outputManagerConfig) { if cfg.channelBufferSize > 0 { t.channelBufferSize = cfg.channelBufferSize } if cfg.flushTimeout > 0 { t.flushTimeout = cfg.flushTimeout } if cfg.flushPollInterval > 0 { t.flushPollInterval = cfg.flushPollInterval } if cfg.readRetryInterval > 0 { t.readRetryInterval = cfg.readRetryInterval } if cfg.eofAckQuietPeriod > 0 { t.eofAckQuietPeriod = cfg.eofAckQuietPeriod } } func (t *outputManager) resolvedChannelBufferSize() int { if t.channelBufferSize > 0 { return t.channelBufferSize } return defaultOutputChannelBufferSize } func (t *outputManager) resolvedFlushTimeout() time.Duration { if t.flushTimeout > 0 { return t.flushTimeout } return defaultOutputFlushTimeout } func (t *outputManager) resolvedFlushPollInterval() time.Duration { if t.flushPollInterval > 0 { return t.flushPollInterval } return defaultOutputFlushPollInterval } func (t *outputManager) resolvedReadRetryInterval() time.Duration { if t.readRetryInterval > 0 { return t.readRetryInterval } return defaultOutputReadRetryInterval } func (t *outputManager) resolvedEOFAckQuietPeriod() time.Duration { if t.eofAckQuietPeriod > 0 { return t.eofAckQuietPeriod } return defaultOutputEOFAckQuietPeriod } // enable atomically switches output mode on. It returns true when it // transitioned from disabled to enabled and false when output mode was already // active. Fresh EOF/EOF-ack channels are created on the off->on transition, // so a concurrent (or repeated) enable can never yank a live EOF channel out // from under an in-flight batch. The lines channel is created once and then // reused for the lifetime of the session. // // One already-enabled case still refreshes the handshake channels: when the // previous batch signaled EOF but the reader never acknowledged it (e.g. // WaitForOutputEOFAck timed out on a slow client), t.eof is already closed. // A new batch inheriting that closed channel would be disabled by // maybeAckEOFLocked as soon as the lines channel is briefly empty, stranding // the batch's remaining output. Refreshing is safe: the old eofAck is closed // first, so a goroutine still blocked on it (e.g. the previous batch mid // quiet-period) is released immediately instead of stalling until its // timeout — its data was already flushed before it signaled EOF, and the // reader can never acknowledge a replaced handshake anyway. // // Every enable() call — transition, join-while-enabled, or stale refresh — // bumps the handshake epoch; see signalEOF. func (t *outputManager) enable() bool { t.mu.Lock() defer t.mu.Unlock() t.epoch++ if t.mode { if t.eofSignaledLocked() { // Stale handshake from an unacknowledged previous batch: release // its waiter and start a fresh handshake so the new batch is not // EOF'd prematurely. t.signalEOFAckLocked() t.resetEOFHandshakeLocked() } return false } t.mode = true if t.lines == nil { t.lines = make(chan []byte, t.resolvedChannelBufferSize()) } // New batch of files: new EOF handshake channels. t.resetEOFHandshakeLocked() return true } // currentEpoch returns the handshake epoch to be captured by a command that // is about to decide whether its batch is over. Capture it BEFORE checking // the pending-work count: joiners increment the pending count before calling // enable(), so a joiner that is invisible to a pending==0 check is guaranteed // to bump the epoch after the capture, invalidating the stale signal. func (t *outputManager) currentEpoch() uint64 { t.mu.Lock() defer t.mu.Unlock() return t.epoch } // eofSignaledLocked reports whether the current EOF channel exists and has // already been closed by signalEOF. The caller must hold t.mu. func (t *outputManager) eofSignaledLocked() bool { if t.eof == nil { return false } select { case <-t.eof: return true default: return false } } // resetEOFHandshakeLocked mints fresh EOF handshake channels for a new batch // and clears the quiet-period tracking. The caller must hold t.mu. func (t *outputManager) resetEOFHandshakeLocked() { t.eof = make(chan struct{}) t.eofAck = make(chan struct{}) t.eofEmptySince = time.Time{} } func (t *outputManager) enabled() bool { t.mu.Lock() defer t.mu.Unlock() return t.mode } func (t *outputManager) hasEOF() bool { t.mu.Lock() defer t.mu.Unlock() return t.eof != nil } // signalEOF closes the EOF channel once, but only when the given epoch // (captured via currentEpoch before the caller's pending-work check) is still // current. If another command joined the output session in between — bumping // the epoch via enable() — the signal is a stale "batch over" decision that // predates the newcomer's output, and closing the (possibly shared) EOF // channel would let the reader disable output mode mid-batch. Such stale // signals are therefore dropped; the newcomer signals EOF itself when its // batch drains. The stale signaler's WaitForOutputEOFAck then waits on the // newcomer's handshake — released when that batch completes, or by timeout. // Either way its own data was already flushed before it signaled. // // Joiners that never signal EOF themselves (output-aggregate/dmap commands, // which are excluded from the cat/grep/tail EOF epilogue) also bump the // epoch. A concurrently finishing cat's signal is then dropped and its ack // wait deterministically times out. That degradation is bounded (one ack // timeout) and loses no data: the cat's output was flushed before it // signaled, and session shutdown flushes whatever remains. func (t *outputManager) signalEOF(epoch uint64) { t.mu.Lock() defer t.mu.Unlock() if epoch != t.epoch { // A newer batch joined after the caller captured its epoch. return } if t.eof == nil { return } select { case <-t.eof: // Already closed default: close(t.eof) } } // signalEOFAckLocked closes the EOF-ack channel once. The caller must hold // t.mu; callers are maybeAckEOFLocked (normal reader-side ack) and enable // (releasing the previous batch's waiter when refreshing a stale handshake). func (t *outputManager) signalEOFAckLocked() { if t.eofAck == nil { return } select { case <-t.eofAck: // Already closed. default: close(t.eofAck) } } // waitForEOFAck blocks until the reader goroutine acknowledges the EOF or the // timeout expires. The ack channel is snapshotted under the lock and waited on // outside it, so a blocked waiter never stalls the reader that has to deliver // the ack. // // A timeout <= 0 is clamped to the default (mirroring the fallback that // OutputEOFAckTimeout uses) as defense-in-depth: every current handshake // replacement leaves the old ack channel closed (a completed handshake was // acked by the reader, and the stale refresh in enable() closes it // explicitly), so a forever-wait cannot presently hang — but the clamp keeps // that true for any future replacement path or caller. func (t *outputManager) waitForEOFAck(timeout time.Duration) bool { t.mu.Lock() eofAck := t.eofAck t.mu.Unlock() if eofAck == nil { return true } if timeout <= 0 { timeout = defaultOutputEOFAckTimeout } timer := time.NewTimer(timeout) defer timer.Stop() select { case <-eofAck: return true case <-timer.C: return false } } func (t *outputManager) channel() chan []byte { t.mu.Lock() defer t.mu.Unlock() return t.lines } func (t *outputManager) channelLen() int { t.mu.Lock() lines := t.lines t.mu.Unlock() if lines == nil { return 0 } return len(lines) } // flush waits until the output lines channel drains or the flush timeout hits. // It polls on a snapshot of the channel taken under the lock (the channel is // never replaced once created), so the reader goroutine is never blocked by a // flusher holding the lock across sleeps. func (t *outputManager) flush(user *user.User) { t.mu.Lock() lines := t.lines t.mu.Unlock() if lines == nil { return } dlog.Server.Debug(user, "Flushing output data", "channelLen", len(lines)) timeout := time.After(t.resolvedFlushTimeout()) for { select { case <-timeout: dlog.Server.Warn(user, "Timeout while flushing output data", "remaining", len(lines)) return default: if len(lines) == 0 { dlog.Server.Debug(user, "Output channel drained successfully") return } // Give the reader time to process. time.Sleep(t.resolvedFlushPollInterval()) } } } // tryRead tries to serve data from output state and channels. // Returns handled=false when caller should continue with normal path. // // It runs on the session output goroutine and holds t.mu for its entire // duration except during the short retry sleep, so command goroutines calling // enable/signalEOF/channelLen never observe torn state (e.g. mode=true with // uninitialized channels). // // Lock ordering: the shouldDropGeneration callback is invoked (via // consumeLocked) while t.mu is held and itself acquires sessionState.mu // (sessionCommandState.currentGeneration), establishing the ordering // output.mu -> sessionState.mu. Nothing may call into outputManager while // holding sessionState.mu, or it would deadlock. func (t *outputManager) tryRead(p []byte, user *user.User, shouldDropGeneration func(uint64) bool) (n int, handled bool) { // tryRead runs on the session output goroutine once per Read (i.e. per output // payload / ~64KB in server mode). Decide trace state once, before taking // the lock, so none of the per-read diagnostics below box their int/string // args or build a []interface{} when trace is off (the default). This also // shortens the t.mu hold time. Locking semantics are unchanged: the guard is // a pure branch and touches no lock. maxLevel is fixed at logger construction. traceEnabled := dlog.Server.TraceEnabled() t.mu.Lock() defer t.mu.Unlock() // Drain buffered remainder data first, regardless of mode: it belongs to a // payload that was already accepted for delivery. This is defensive — with // the current code the combination buffer-nonempty + mode-off cannot arise // (only maybeAckEOFLocked clears mode, and it runs only after the buffer // has been drained) — but draining first keeps delivery correct should // that invariant ever change. if len(t.buffer) > 0 { if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "using buffered output data", "bufferedLen", len(t.buffer)) } n = copy(p, t.buffer) t.buffer = t.buffer[n:] if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "after buffer read", "copied", n, "remaining", len(t.buffer)) } return n, true } if !t.mode { return 0, false } if t.lines == nil { return 0, false } if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "checking outputLines channel", "channelLen", len(t.lines)) } for { select { case outputData := <-t.lines: if n, delivered := t.consumeLocked(p, outputData, user, traceEnabled, shouldDropGeneration); delivered { return n, true } continue default: } // Recompute per iteration: after draining stale-generation entries the // channel may be empty, in which case the retry sleep is pointless. if len(t.lines) > 0 { if outputData, received := t.retryReceiveLocked(user, traceEnabled); received { if n, delivered := t.consumeLocked(p, outputData, user, traceEnabled, shouldDropGeneration); delivered { if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "got data after wait") } return n, true } continue } } t.maybeAckEOFLocked(user) if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "no data in outputLines, falling through") } return 0, false } } // retryReceiveLocked waits one retry interval for slow producers and then // attempts a non-blocking receive from the lines channel. The lock is released // during the sleep so command goroutines (enable, signalEOF, flush, ...) are // not stalled by the reader's retry backoff; the lines channel is never // replaced once created, so re-checking it after re-locking is safe. The // caller must hold t.mu; it is held again on return. func (t *outputManager) retryReceiveLocked(user *user.User, traceEnabled bool) (outputData []byte, received bool) { if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "channel has data but not available, waiting") } retryInterval := t.resolvedReadRetryInterval() t.mu.Unlock() time.Sleep(retryInterval) t.mu.Lock() select { case outputData = <-t.lines: return outputData, true default: // Still no data. return nil, false } } // consumeLocked decodes a output payload, drops it when its generation is // stale, and otherwise copies it into p, buffering any remainder for the next // read. Returns delivered=false when the payload was dropped. The caller must // hold t.mu. func (t *outputManager) consumeLocked(p, outputData []byte, user *user.User, traceEnabled bool, shouldDropGeneration func(uint64) bool) (n int, delivered bool) { generation, decodedData := decodeGeneratedBytes(outputData) if shouldDropGeneration != nil && shouldDropGeneration(generation) { t.eofEmptySince = time.Time{} return 0, false } if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "got data from outputLines", "dataLen", len(decodedData)) } t.eofEmptySince = time.Time{} n = copy(p, decodedData) if n < len(decodedData) { t.buffer = decodedData[n:] if traceEnabled { dlog.Server.Trace(user, "baseHandler.Read", "buffering remaining data", "bufferedLen", len(t.buffer)) } } return n, true } // maybeAckEOFLocked disables output mode and acknowledges the EOF once EOF has // been signaled and the lines channel has stayed empty for the quiet period. // The caller must hold t.mu. func (t *outputManager) maybeAckEOFLocked(user *user.User) { if t.eof == nil { return } select { case <-t.eof: default: return } if len(t.lines) > 0 { t.eofEmptySince = time.Time{} return } if t.eofEmptySince.IsZero() { t.eofEmptySince = time.Now() return } if time.Since(t.eofEmptySince) >= t.resolvedEOFAckQuietPeriod() { dlog.Server.Trace(user, "baseHandler.Read", "EOF acknowledged and channel stable-empty, disabling output mode") t.mode = false t.signalEOFAckLocked() } }