diff --git a/transport/internet/splithttp/client.go b/transport/internet/splithttp/client.go index a46c2b8ca..a2b5f75e8 100644 --- a/transport/internet/splithttp/client.go +++ b/transport/internet/splithttp/client.go @@ -68,7 +68,7 @@ func (c *DefaultDialerClient) OpenStream(ctx context.Context, url string, sessio } c.transportConfig.FillStreamRequest(req, sessionId, "") - wrc = &WaitReadCloser{Wait: make(chan struct{})} + wrc = &WaitReadCloser{wait: done.New()} go func() { resp, err := c.client.Do(req) if err != nil { @@ -188,38 +188,35 @@ func (c *DefaultDialerClient) Close() error { } type WaitReadCloser struct { - Wait chan struct{} - io.ReadCloser + wait *done.Instance + reader atomic.Pointer[io.ReadCloser] } func (w *WaitReadCloser) Set(rc io.ReadCloser) { - w.ReadCloser = rc - defer func() { - if recover() != nil { - rc.Close() + w.reader.Store(&rc) + if w.wait.Done() { + if p := w.reader.Swap(nil); p != nil { + (*p).Close() } - }() - close(w.Wait) + } + w.wait.Close() } func (w *WaitReadCloser) Read(b []byte) (int, error) { - if w.ReadCloser == nil { - if <-w.Wait; w.ReadCloser == nil { + rc := w.reader.Load() + if rc == nil { + <-w.wait.Wait() + if rc = w.reader.Load(); rc == nil { return 0, io.ErrClosedPipe } } - return w.ReadCloser.Read(b) + return (*rc).Read(b) } func (w *WaitReadCloser) Close() error { - if w.ReadCloser != nil { - return w.ReadCloser.Close() + w.wait.Close() + if p := w.reader.Swap(nil); p != nil { + return (*p).Close() } - defer func() { - if recover() != nil && w.ReadCloser != nil { - w.ReadCloser.Close() - } - }() - close(w.Wait) return nil }