mirror of
https://github.com/HDT3213/godis.git
synced 2025-10-05 00:42:43 +08:00
159 lines
3.0 KiB
Go
159 lines
3.0 KiB
Go
package connection
|
|
|
|
import (
|
|
"bytes"
|
|
"github.com/hdt3213/godis/lib/sync/wait"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Connection represents a connection with a redis-cli
|
|
type Connection struct {
|
|
conn net.Conn
|
|
|
|
// waiting until reply finished
|
|
waitingReply wait.Wait
|
|
|
|
// lock while server sending response
|
|
mu sync.Mutex
|
|
|
|
// subscribing channels
|
|
subs map[string]bool
|
|
|
|
// password may be changed by CONFIG command during runtime, so store the password
|
|
password string
|
|
|
|
// queued commands for `multi`
|
|
multiState bool
|
|
queue [][][]byte
|
|
}
|
|
|
|
// RemoteAddr returns the remote network address
|
|
func (c *Connection) RemoteAddr() net.Addr {
|
|
return c.conn.RemoteAddr()
|
|
}
|
|
|
|
// Close disconnect with the client
|
|
func (c *Connection) Close() error {
|
|
c.waitingReply.WaitWithTimeout(10 * time.Second)
|
|
_ = c.conn.Close()
|
|
return nil
|
|
}
|
|
|
|
// NewConn creates Connection instance
|
|
func NewConn(conn net.Conn) *Connection {
|
|
return &Connection{
|
|
conn: conn,
|
|
}
|
|
}
|
|
|
|
// Write sends response to client over tcp connection
|
|
func (c *Connection) Write(b []byte) error {
|
|
if len(b) == 0 {
|
|
return nil
|
|
}
|
|
c.mu.Lock()
|
|
c.waitingReply.Add(1)
|
|
defer func() {
|
|
c.waitingReply.Done()
|
|
c.mu.Unlock()
|
|
}()
|
|
|
|
_, err := c.conn.Write(b)
|
|
return err
|
|
}
|
|
|
|
// Subscribe add current connection into subscribers of the given channel
|
|
func (c *Connection) Subscribe(channel string) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
if c.subs == nil {
|
|
c.subs = make(map[string]bool)
|
|
}
|
|
c.subs[channel] = true
|
|
}
|
|
|
|
// UnSubscribe removes current connection into subscribers of the given channel
|
|
func (c *Connection) UnSubscribe(channel string) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
if len(c.subs) == 0 {
|
|
return
|
|
}
|
|
delete(c.subs, channel)
|
|
}
|
|
|
|
// SubsCount returns the number of subscribing channels
|
|
func (c *Connection) SubsCount() int {
|
|
return len(c.subs)
|
|
}
|
|
|
|
// GetChannels returns all subscribing channels
|
|
func (c *Connection) GetChannels() []string {
|
|
if c.subs == nil {
|
|
return make([]string, 0)
|
|
}
|
|
channels := make([]string, len(c.subs))
|
|
i := 0
|
|
for channel := range c.subs {
|
|
channels[i] = channel
|
|
i++
|
|
}
|
|
return channels
|
|
}
|
|
|
|
// SetPassword stores password for authentication
|
|
func (c *Connection) SetPassword(password string) {
|
|
c.password = password
|
|
}
|
|
|
|
// GetPassword get password for authentication
|
|
func (c *Connection) GetPassword() string {
|
|
return c.password
|
|
}
|
|
|
|
func (c *Connection) InMultiState() bool {
|
|
return c.multiState
|
|
}
|
|
|
|
func (c *Connection) SetMultiState(state bool) {
|
|
c.multiState = state
|
|
}
|
|
|
|
func (c *Connection) GetQueuedCmdLine() [][][]byte {
|
|
return c.queue
|
|
}
|
|
|
|
func (c *Connection) EnqueueCmd(cmdLine [][]byte) {
|
|
c.queue = append(c.queue, cmdLine)
|
|
}
|
|
|
|
func (c *Connection) ClearQueuedCmds() {
|
|
c.queue = nil
|
|
}
|
|
|
|
// FakeConn implements redis.Connection for test
|
|
type FakeConn struct {
|
|
Connection
|
|
buf bytes.Buffer
|
|
}
|
|
|
|
// Write writes data to buffer
|
|
func (c *FakeConn) Write(b []byte) error {
|
|
c.buf.Write(b)
|
|
return nil
|
|
}
|
|
|
|
// Clean resets the buffer
|
|
func (c *FakeConn) Clean() {
|
|
c.buf.Reset()
|
|
}
|
|
|
|
// Bytes returns written data
|
|
func (c *FakeConn) Bytes() []byte {
|
|
return c.buf.Bytes()
|
|
}
|