Skip to content

Commit

Permalink
Merge branch 'v4' into fix-sentinel_auth
Browse files Browse the repository at this point in the history
  • Loading branch information
EquentR authored Nov 25, 2024
2 parents e92b8de + 5da4fda commit 53f5ddb
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 2 deletions.
58 changes: 58 additions & 0 deletions internal/client/redis.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import (
"regexp"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"

"RedisShake/internal/client/proto"
Expand All @@ -20,6 +22,9 @@ type Redis struct {
writer *bufio.Writer
protoReader *proto.Reader
protoWriter *proto.Writer
timer *time.Timer
sendCount uint64
mu sync.Mutex
}

func NewSentinelMasterClient(ctx context.Context, address string, username string, password string, Tls bool) *Redis {
Expand Down Expand Up @@ -81,6 +86,9 @@ func NewRedisClient(ctx context.Context, address string, username string, passwo
r = NewRedisClient(ctx, replicaInfo.BestReplica, username, password, Tls, false)
}

r.timer = time.NewTimer(time.Second)
go r.autoFlush(ctx)

return r
}

Expand Down Expand Up @@ -185,11 +193,54 @@ func (r *Redis) SendBytes(buf []byte) {
r.flush()
}

func (r *Redis) SendBytesBuff(buf []byte) {
r.mu.Lock()
defer r.mu.Unlock()
_, err := r.writer.Write(buf)
if err != nil {
log.Panicf(err.Error())
}
r.flushBuff()
}

func (r *Redis) flushBuff() {
if atomic.AddUint64(&r.sendCount, 1)%100 != 0 {
return
}
if !r.timer.Stop() {
select {
case <-r.timer.C:
default:
}
}
r.timer.Reset(time.Second)
r.flush()
}

func (r *Redis) flush() {
err := r.writer.Flush()
if err != nil {
log.Panicf(err.Error())
}
atomic.StoreUint64(&r.sendCount, 0)
}

func (r *Redis) autoFlush(ctx context.Context) {
for {
select {
case <-ctx.Done():
return
case <-r.timer.C:
if atomic.LoadUint64(&r.sendCount) > 0 {
r.mu.Lock()
err := r.writer.Flush()
r.mu.Unlock()
if err != nil {
log.Panicf(err.Error())
}
}
}
}
}

func (r *Redis) Receive() (interface{}, error) {
Expand Down Expand Up @@ -217,6 +268,13 @@ func (r *Redis) Close() {
if err := r.conn.Close(); err != nil {
log.Infof("close redis conn err: %s\n", err.Error())
}
// release the timer
if !r.timer.Stop() {
select {
case <-r.timer.C:
default:
}
}
}

/* Commands */
Expand Down
2 changes: 1 addition & 1 deletion internal/rdb/types/stream.go
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ func (o *StreamObject) readStream() {
cmdC <- []string{"XGROUP", "CREATE", masterKey, groupName, lastid}

/* Load group offset. */
if typeByte == rdbTypeStreamListpacks2 {
if typeByte >= rdbTypeStreamListpacks2 {
_ = structure.ReadLength(rd) // offset
}

Expand Down
11 changes: 10 additions & 1 deletion internal/writer/redis_standalone_writer.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ type RedisWriterOptions struct {
SentinelPassword string `mapstructure:"sentinel_password" default:""`
Tls bool `mapstructure:"tls" default:"false"`
OffReply bool `mapstructure:"off_reply" default:"false"`
BuffSend bool `mapstructure:"buff_send" default:"false"`
}

type redisStandaloneWriter struct {
Expand All @@ -41,6 +42,8 @@ type redisStandaloneWriter struct {
ch chan *entry.Entry
chWg sync.WaitGroup

buffSend bool

stat struct {
Name string `json:"name"`
UnansweredBytes int64 `json:"unanswered_bytes"`
Expand All @@ -54,6 +57,7 @@ func NewRedisStandaloneWriter(ctx context.Context, opts *RedisWriterOptions) Wri
rw.stat.Name = "writer_" + strings.Replace(opts.Address, ":", "_", -1)
rw.client = client.NewRedisClient(ctx, opts.Address, opts.Username, opts.Password, opts.Tls, false)
rw.ch = make(chan *entry.Entry, 1024)
rw.buffSend = opts.BuffSend
if opts.OffReply {
log.Infof("turn off the reply of write")
rw.offReply = true
Expand Down Expand Up @@ -95,7 +99,12 @@ func (w *redisStandaloneWriter) StartWrite(ctx context.Context) chan *entry.Entr
atomic.AddInt64(&w.stat.UnansweredBytes, e.SerializedSize)
atomic.AddInt64(&w.stat.UnansweredEntries, 1)
}
w.client.SendBytes(bytes)
if w.buffSend {
w.client.SendBytesBuff(bytes)
} else {
w.client.SendBytes(bytes)
}

}
w.chWg.Done()
}()
Expand Down
1 change: 1 addition & 0 deletions shake.toml
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ sentinel_username = "" # keep empty if not using sentinel ACL
sentinel_password = "" # keep empty if sentinel no authentication is required
tls = false
off_reply = false # turn off the server reply
buff_send = false # buffer send, default false. may be a sync delay when true, but it can greatly improve the speed

[filter]
# Allow keys with specific prefixes or suffixes
Expand Down

0 comments on commit 53f5ddb

Please sign in to comment.