fix: add flush timeout for logs (#1007)

Co-authored-by: Luke Lombardi <[email protected]>

pkg/abstractions/common/logs.go

@@ -12,6 +12,10 @@ import (
 	"github.com/beam-cloud/beta9/pkg/types"
 )
 
+const (
+	flushLogsTimeout = 500 * time.Millisecond
+)
+
 type LogStreamOpts struct {
 	SendCallback    func(o common.OutputMsg) error
 	ExitCallback    func(exitCode int32) error
@@ -97,6 +101,25 @@ _stream:
 				exitCode = -1
 			}
 
+			// After the container exits, flush remaining logs for up to flushLogsTimeout milliseconds
+			flushLogsTimer := time.NewTimer(flushLogsTimeout)
+			defer flushLogsTimer.Stop()
+
+		_flush:
+			for {
+				select {
+				case o, ok := <-outputChan:
+					if !ok {
+						break _flush
+					}
+					if err := l.sendCallback(o); err != nil {
+						break _flush
+					}
+				case <-flushLogsTimer.C:
+					break _flush
+				}
+			}
+
 			if err := l.exitCallback(int32(exitCode)); err != nil {
 				break _stream
 			}