/
conn.go
129 lines (103 loc) · 2.08 KB
/
conn.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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package pool
import (
"context"
"net"
"strconv"
"sync/atomic"
"time"
"github.com/go-pg/pg/v9/internal"
)
var noDeadline = time.Time{}
type Conn struct {
netConn net.Conn
rd *internal.BufReader
ProcessID int32
SecretKey int32
lastID int64
createdAt time.Time
usedAt uint32 // atomic
pooled bool
Inited bool
}
func NewConn(netConn net.Conn) *Conn {
cn := &Conn{
rd: internal.NewBufReader(netConn),
createdAt: time.Now(),
}
cn.SetNetConn(netConn)
cn.SetUsedAt(time.Now())
return cn
}
func (cn *Conn) UsedAt() time.Time {
unix := atomic.LoadUint32(&cn.usedAt)
return time.Unix(int64(unix), 0)
}
func (cn *Conn) SetUsedAt(tm time.Time) {
atomic.StoreUint32(&cn.usedAt, uint32(tm.Unix()))
}
func (cn *Conn) RemoteAddr() net.Addr {
return cn.netConn.RemoteAddr()
}
func (cn *Conn) SetNetConn(netConn net.Conn) {
cn.netConn = netConn
cn.rd.Reset(netConn)
}
func (cn *Conn) NetConn() net.Conn {
return cn.netConn
}
func (cn *Conn) NextID() string {
cn.lastID++
return strconv.FormatInt(cn.lastID, 10)
}
func (cn *Conn) WithReader(
ctx context.Context, timeout time.Duration, fn func(rd *internal.BufReader) error,
) error {
err := cn.netConn.SetReadDeadline(cn.deadline(ctx, timeout))
if err != nil {
return err
}
return fn(cn.rd)
}
func (cn *Conn) WithWriter(
ctx context.Context, timeout time.Duration, fn func(wb *WriteBuffer) error,
) error {
err := cn.netConn.SetWriteDeadline(cn.deadline(ctx, timeout))
if err != nil {
return err
}
wb := getWriteBuffer()
defer putWriteBuffer(wb)
wb.Reset()
err = fn(wb)
if err != nil {
return err
}
_, err = cn.netConn.Write(wb.Bytes)
return err
}
func (cn *Conn) Close() error {
return cn.netConn.Close()
}
func (cn *Conn) deadline(ctx context.Context, timeout time.Duration) time.Time {
tm := time.Now()
cn.SetUsedAt(tm)
if timeout > 0 {
tm = tm.Add(timeout)
}
if ctx != nil {
deadline, ok := ctx.Deadline()
if ok {
if timeout == 0 {
return deadline
}
if deadline.Before(tm) {
return deadline
}
return tm
}
}
if timeout > 0 {
return tm
}
return noDeadline
}