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() }