Something went wrong. Try again.
small gleam coding and (not yet) persistent agent daemon with a detachable cli forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104package daemon
import ( "io" "net/http" "strings" "sync"
"github.com/klauspost/compress/zstd")
// Request bodies at least this large are compressed when the daemon// advertises zstd_requests; smaller ones are not worth a frame.const compressedRequestBytes = 16 * 1024
// The window a frame declares is the memory the decoder must hold for it, and// below 2 MiB klauspost allocates twice that: the daemon compresses with a// 64 KiB window (albedo_http_api's zstd_params), so the decoder a session's// stream keeps for the life of the client costs about 200 KiB. This cap only// refuses a frame no daemon of ours sends; level 1's default window is// 512 KiB, which is what a daemon that predates the 64 KiB window declares.const zstdWindowBytes = 8 << 20
// EncodeAll is safe for concurrent use; one encoder amortises its tables// across every request body.var zstdEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedFastest), zstd.WithEncoderConcurrency(1))
// Decoders are the expensive half: a fresh one allocates its frame state per// response. Each body borrows one for its exclusive use and hands it back// when it closes.var zstdDecoders = sync.Pool{New: func() any { decoder, _ := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxWindow(zstdWindowBytes)) return decoder}}
// requestBody returns the payload to send and its Content-Encoding, empty// for identity.func requestBody(conn *Connection, payload []byte) ([]byte, string) { if len(payload) < compressedRequestBytes || conn.capability("zstd_requests") < 1 { return payload, "" } return zstdEncoder.EncodeAll(payload, nil), "zstd"}
// codingTransport asks for zstd responses and decodes them, the way// net/http does for gzip when no Accept-Encoding is set.type codingTransport struct{ base http.RoundTripper }
func (transport codingTransport) RoundTrip(req *http.Request) (*http.Response, error) { if req.Header.Get("Accept-Encoding") == "" { req = req.Clone(req.Context()) req.Header.Set("Accept-Encoding", "zstd") } res, err := transport.base.RoundTrip(req) if err != nil || !strings.EqualFold(res.Header.Get("Content-Encoding"), "zstd") { return res, err } decoder := zstdDecoders.Get().(*zstd.Decoder) // A frame that is not zstd fails on the first Read, as it does for any // streamed decoder; only a closed decoder fails here. if err := decoder.Reset(res.Body); err != nil { decoder.Close() _ = res.Body.Close() return nil, err } res.Body = &zstdBody{decoder: decoder, body: res.Body} res.Header.Del("Content-Encoding") res.Header.Del("Content-Length") res.ContentLength = -1 res.Uncompressed = true return res, nil}
func (transport codingTransport) CloseIdleConnections() { if closer, ok := transport.base.(interface{ CloseIdleConnections() }); ok { closer.CloseIdleConnections() }}
type zstdBody struct { decoder *zstd.Decoder body io.ReadCloser released bool}
func (body *zstdBody) Read(target []byte) (int, error) { if body.released { return 0, io.ErrClosedPipe } return body.decoder.Read(target)}
func (body *zstdBody) Close() error { if body.released { return nil } body.released = true err := body.body.Close() // Detach the closed body before the decoder is reused for another stream. _ = body.decoder.Reset(nil) zstdDecoders.Put(body.decoder) return err}