-
Notifications
You must be signed in to change notification settings - Fork 338
/
stream.go
74 lines (62 loc) · 1.65 KB
/
stream.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
// Copyright 2020 The Swarm Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package libp2p
import (
"errors"
"io"
"time"
"github.com/ethersphere/bee/v2/pkg/p2p"
"github.com/libp2p/go-libp2p/core/network"
)
var (
closeDeadline = 30 * time.Second
errExpectedEof = errors.New("read: expected eof")
)
var _ p2p.Stream = (*stream)(nil)
type stream struct {
network.Stream
headers map[string][]byte
responseHeaders map[string][]byte
metrics metrics
}
func newStream(s network.Stream, metrics metrics) *stream {
return &stream{Stream: s, metrics: metrics}
}
func (s *stream) Headers() p2p.Headers {
return s.headers
}
func (s *stream) ResponseHeaders() p2p.Headers {
return s.responseHeaders
}
func (s *stream) Reset() error {
defer s.metrics.StreamResetCount.Inc()
return s.Stream.Reset()
}
func (s *stream) FullClose() error {
defer s.metrics.ClosedStreamCount.Inc()
// close the stream to make sure it is gc'd
defer s.Close()
if err := s.CloseWrite(); err != nil {
_ = s.Stream.Reset()
return err
}
// So we don't wait forever
_ = s.SetDeadline(time.Now().Add(closeDeadline))
// We *have* to observe the EOF. Otherwise, we leak the stream.
// Now, technically, we should do this *before*
// returning from SendMessage as the message
// hasn't really been sent yet until we see the
// EOF but we don't actually *know* what
// protocol the other side is speaking.
n, err := s.Read([]byte{0})
if n > 0 || err == nil {
_ = s.Stream.Reset()
return errExpectedEof
}
if !errors.Is(err, io.EOF) {
_ = s.Stream.Reset()
return err
}
return nil
}