package notify import ( "b612.me/stario" "context" "errors" "io" "net" "strings" "sync" "time" ) var transportConnWriteGates sync.Map var errTransportFrameQueueUnavailable = errors.New("transport frame queue is unavailable") type connWriteGateRef struct { mu sync.Mutex gate chan struct{} refs int } type vectoredBuffersWriter interface { WriteBuffers(*net.Buffers) (int64, error) } type vectoredConnUnwrapper interface { UnwrapConn() net.Conn } func writeFullToConn(conn net.Conn, data []byte) error { if conn == nil { return net.ErrClosed } return withRawConnWriteLock(conn, func(conn net.Conn) error { return writeFullToConnUnlocked(conn, data) }) } func writeFullToConnUnlocked(conn net.Conn, data []byte) error { if conn == nil { return net.ErrClosed } return writeFullToWriterUnlocked(conn, data) } func writeFullToWriterUnlocked(writer io.Writer, data []byte) error { if writer == nil { return io.ErrClosedPipe } for len(data) > 0 { n, err := writer.Write(data) if n > 0 { data = data[n:] } if err != nil { return err } if n == 0 { return io.ErrNoProgress } } return nil } func writeNetBuffersFullUnlocked(conn net.Conn, buffers net.Buffers) error { if conn == nil { return net.ErrClosed } writer, writeFn := vectoredWriteStrategy(conn) if writeFn == nil { return writeRemainingBuffersUnlocked(conn, buffers) } n, err := writeFn(&buffers) if err != nil { return err } if len(buffers) == 0 { return nil } if n == 0 { return io.ErrNoProgress } return writeRemainingBuffersUnlocked(writer, buffers) } func vectoredWriteStrategy(conn net.Conn) (io.Writer, func(*net.Buffers) (int64, error)) { current := conn for depth := 0; depth < 8 && current != nil; depth++ { if writer, ok := current.(vectoredBuffersWriter); ok { target := current return target, writer.WriteBuffers } switch target := current.(type) { case *net.TCPConn: return target, func(bufs *net.Buffers) (int64, error) { return bufs.WriteTo(target) } case *net.UnixConn: return target, func(bufs *net.Buffers) (int64, error) { return bufs.WriteTo(target) } } unwrapper, ok := current.(vectoredConnUnwrapper) if !ok { break } next := unwrapper.UnwrapConn() if next == nil || next == current { break } current = next } return nil, nil } func writeRemainingBuffersUnlocked(writer io.Writer, buffers net.Buffers) error { for _, part := range buffers { if len(part) == 0 { continue } if err := writeFullToWriterUnlocked(writer, part); err != nil { return err } } return nil } func withRawConnWriteLock(conn net.Conn, fn func(net.Conn) error) error { return withRawConnWriteLockDeadline(conn, time.Time{}, fn) } func withRawConnWriteLockDeadline(conn net.Conn, deadline time.Time, fn func(net.Conn) error) error { _, err := withRawConnWriteLockContextDeadline(context.Background(), conn, deadline, fn) return err } func withRawConnWriteLockContextDeadline(ctx context.Context, conn net.Conn, deadline time.Time, fn func(net.Conn) error) (bool, error) { if conn == nil { return false, net.ErrClosed } if ctx == nil { ctx = context.Background() } gateRef := retainRawConnWriteGate(conn) defer releaseRawConnWriteGate(conn, gateRef) gate := gateRef.gate if err := lockWriteGateContextDeadline(ctx, nil, gate, deadline); err != nil { return false, err } defer func() { gate <- struct{}{} }() if err := ctx.Err(); err != nil { return false, err } deadline = earlierWriteDeadline(deadline, contextDeadline(ctx)) if !deadline.IsZero() { if err := conn.SetWriteDeadline(deadline); err != nil { return true, err } defer func() { _ = conn.SetWriteDeadline(time.Time{}) }() } return true, fn(conn) } func retainRawConnWriteGate(conn net.Conn) *connWriteGateRef { if conn == nil { return &connWriteGateRef{gate: newConnWriteGate(), refs: 1} } for { candidate := &connWriteGateRef{gate: newConnWriteGate(), refs: 1} actual, loaded := transportConnWriteGates.LoadOrStore(conn, candidate) if !loaded { return candidate } ref := actual.(*connWriteGateRef) ref.mu.Lock() if ref.refs > 0 { ref.refs++ ref.mu.Unlock() return ref } ref.mu.Unlock() transportConnWriteGates.CompareAndDelete(conn, ref) } } func releaseRawConnWriteGate(conn net.Conn, ref *connWriteGateRef) { if ref == nil { return } ref.mu.Lock() if ref.refs > 0 { ref.refs-- } remove := ref.refs == 0 ref.mu.Unlock() if remove && conn != nil { transportConnWriteGates.CompareAndDelete(conn, ref) } } func newConnWriteGate() chan struct{} { gate := make(chan struct{}, 1) gate <- struct{}{} return gate } func lockWriteGateContextDeadline(ctx context.Context, stop <-chan struct{}, gate chan struct{}, deadline time.Time) error { if ctx == nil { ctx = context.Background() } select { case <-ctx.Done(): return ctx.Err() case <-stop: return net.ErrClosed default: } select { case <-gate: return nil default: } if deadline.IsZero() { select { case <-ctx.Done(): return ctx.Err() case <-stop: return net.ErrClosed case <-gate: return nil } } wait := time.Until(deadline) if wait <= 0 { return context.DeadlineExceeded } timer := time.NewTimer(wait) defer timer.Stop() select { case <-ctx.Done(): return ctx.Err() case <-stop: return net.ErrClosed case <-timer.C: return context.DeadlineExceeded case <-gate: return nil } } func shorterPositiveDuration(left time.Duration, right time.Duration) time.Duration { left = maxDuration(0, left) right = maxDuration(0, right) if left == 0 { return right } if right == 0 || left < right { return left } return right } func writeFramedPayloadUnlocked(conn net.Conn, queue *stario.StarQueue, payload []byte) error { if conn == nil { return net.ErrClosed } if queue == nil { return errTransportFrameQueueUnavailable } if isPacketTransportConn(conn) { return writeFullToConnUnlocked(conn, queue.BuildMessage(payload)) } return queue.WriteFrameBuffers(conn, payload) } func writeFramedPayloadBatchUnlocked(conn net.Conn, queue *stario.StarQueue, payloads [][]byte) error { if conn == nil { return net.ErrClosed } if queue == nil { return errTransportFrameQueueUnavailable } if len(payloads) == 0 { return nil } if isPacketTransportConn(conn) { for _, payload := range payloads { if err := writeFullToConnUnlocked(conn, queue.BuildMessage(payload)); err != nil { return err } } return nil } return queue.WriteFramesBuffers(conn, payloads...) } func isPacketTransportConn(conn net.Conn) bool { if conn == nil { return false } if _, ok := conn.(*net.UDPConn); ok { return true } return isPacketNetwork(addrNetwork(conn.LocalAddr())) || isPacketNetwork(addrNetwork(conn.RemoteAddr())) } func addrNetwork(addr net.Addr) string { if addr == nil { return "" } return addr.Network() } func isPacketNetwork(network string) bool { switch strings.ToLower(network) { case "udp", "udp4", "udp6": return true default: return false } }