// Copyright 2011 Gary Burd // // Licensed under the Apache License, Version 2.0 (the "License"): you may // not use this file except in compliance with the License. You may obtain // a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the // License for the specific language governing permissions and limitations // under the License. package redis import ( "context" "errors" "io" "reflect" "sync" "testing" "time" "github.com/go-kratos/kratos/pkg/container/pool" ) type poolTestConn struct { d *poolDialer err error c Conn ctx context.Context } func (c *poolTestConn) Flush() error { return c.c.Flush() } func (c *poolTestConn) Receive() (reply interface{}, err error) { return c.c.Receive() } func (c *poolTestConn) WithContext(ctx context.Context) Conn { c.c.WithContext(ctx) c.ctx = ctx return c } func (c *poolTestConn) Close() error { c.d.mu.Lock() c.d.open-- c.d.mu.Unlock() return c.c.Close() } func (c *poolTestConn) Err() error { return c.err } func (c *poolTestConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) { if commandName == "ERR" { c.err = args[0].(error) commandName = "PING" } if commandName != "" { c.d.commands = append(c.d.commands, commandName) } return c.c.Do(commandName, args...) } func (c *poolTestConn) Send(commandName string, args ...interface{}) error { c.d.commands = append(c.d.commands, commandName) return c.c.Send(commandName, args...) } type poolDialer struct { mu sync.Mutex t *testing.T dialed int open int commands []string dialErr error } func (d *poolDialer) dial() (Conn, error) { d.mu.Lock() d.dialed += 1 dialErr := d.dialErr d.mu.Unlock() if dialErr != nil { return nil, d.dialErr } c, err := DialDefaultServer() if err != nil { return nil, err } d.mu.Lock() d.open += 1 d.mu.Unlock() return &poolTestConn{d: d, c: c}, nil } func (d *poolDialer) check(message string, p *Pool, dialed, open int) { d.mu.Lock() if d.dialed != dialed { d.t.Errorf("%s: dialed=%d, want %d", message, d.dialed, dialed) } if d.open != open { d.t.Errorf("%s: open=%d, want %d", message, d.open, open) } // if active := p.ActiveCount(); active != open { // d.t.Errorf("%s: active=%d, want %d", message, active, open) // } d.mu.Unlock() } func TestPoolReuse(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } var err error for i := 0; i < 10; i++ { c1 := p.Get(context.TODO()) c1.Do("PING") c2 := p.Get(context.TODO()) c2.Do("PING") c1.Close() c2.Close() } d.check("before close", p, 2, 2) err = p.Close() if err != nil { t.Fatal(err) } d.check("after close", p, 2, 0) } func TestPoolMaxIdle(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() for i := 0; i < 10; i++ { c1 := p.Get(context.TODO()) c1.Do("PING") c2 := p.Get(context.TODO()) c2.Do("PING") c3 := p.Get(context.TODO()) c3.Do("PING") c1.Close() c2.Close() c3.Close() } d.check("before close", p, 12, 2) p.Close() d.check("after close", p, 12, 0) } func TestPoolError(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c := p.Get(context.TODO()) c.Do("ERR", io.EOF) if c.Err() == nil { t.Errorf("expected c.Err() != nil") } c.Close() c = p.Get(context.TODO()) c.Do("ERR", io.EOF) c.Close() d.check(".", p, 2, 0) } func TestPoolClose(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c1 := p.Get(context.TODO()) c1.Do("PING") c2 := p.Get(context.TODO()) c2.Do("PING") c3 := p.Get(context.TODO()) c3.Do("PING") c1.Close() if _, err := c1.Do("PING"); err == nil { t.Errorf("expected error after connection closed") } c2.Close() c2.Close() p.Close() d.check("after pool close", p, 3, 1) if _, err := c1.Do("PING"); err == nil { t.Errorf("expected error after connection and pool closed") } c3.Close() d.check("after conn close", p, 3, 0) c1 = p.Get(context.TODO()) if _, err := c1.Do("PING"); err == nil { t.Errorf("expected error after pool closed") } } func TestPoolConcurrenSendReceive(t *testing.T) { p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return DialDefaultServer() } defer p.Close() c := p.Get(context.TODO()) done := make(chan error, 1) go func() { _, err := c.Receive() done <- err }() c.Send("PING") c.Flush() err := <-done if err != nil { t.Fatalf("Receive() returned error %v", err) } _, err = c.Do("") if err != nil { t.Fatalf("Do() returned error %v", err) } c.Close() } func TestPoolMaxActive(t *testing.T) { d := poolDialer{t: t} conf := getTestConfig(testRedisAddr) conf.Config = &pool.Config{ Active: 2, Idle: 2, } p := NewPool(conf) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c1 := p.Get(context.TODO()) c1.Do("PING") c2 := p.Get(context.TODO()) c2.Do("PING") d.check("1", p, 2, 2) c3 := p.Get(context.TODO()) if _, err := c3.Do("PING"); err != pool.ErrPoolExhausted { t.Errorf("expected pool exhausted") } c3.Close() d.check("2", p, 2, 2) c2.Close() d.check("3", p, 2, 2) c3 = p.Get(context.TODO()) if _, err := c3.Do("PING"); err != nil { t.Errorf("expected good channel, err=%v", err) } c3.Close() d.check("4", p, 2, 2) } func TestPoolMonitorCleanup(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c := p.Get(context.TODO()) c.Send("MONITOR") c.Close() d.check("", p, 1, 0) } func TestPoolPubSubCleanup(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c := p.Get(context.TODO()) c.Send("SUBSCRIBE", "x") c.Close() want := []string{"SUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil c = p.Get(context.TODO()) c.Send("PSUBSCRIBE", "x*") c.Close() want = []string{"PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil } func TestPoolTransactionCleanup(t *testing.T) { d := poolDialer{t: t} p := NewPool(testConfig) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c := p.Get(context.TODO()) c.Do("WATCH", "key") c.Do("PING") c.Close() want := []string{"WATCH", "PING", "UNWATCH"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil c = p.Get(context.TODO()) c.Do("WATCH", "key") c.Do("UNWATCH") c.Do("PING") c.Close() want = []string{"WATCH", "UNWATCH", "PING"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil c = p.Get(context.TODO()) c.Do("WATCH", "key") c.Do("MULTI") c.Do("PING") c.Close() want = []string{"WATCH", "MULTI", "PING", "DISCARD"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil c = p.Get(context.TODO()) c.Do("WATCH", "key") c.Do("MULTI") c.Do("DISCARD") c.Do("PING") c.Close() want = []string{"WATCH", "MULTI", "DISCARD", "PING"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil c = p.Get(context.TODO()) c.Do("WATCH", "key") c.Do("MULTI") c.Do("EXEC") c.Do("PING") c.Close() want = []string{"WATCH", "MULTI", "EXEC", "PING"} if !reflect.DeepEqual(d.commands, want) { t.Errorf("got commands %v, want %v", d.commands, want) } d.commands = nil } func startGoroutines(p *Pool, cmd string, args ...interface{}) chan error { errs := make(chan error, 10) for i := 0; i < cap(errs); i++ { go func() { c := p.Get(context.TODO()) _, err := c.Do(cmd, args...) errs <- err c.Close() }() } // Wait for goroutines to block. time.Sleep(time.Second / 4) return errs } func TestWaitPoolDialError(t *testing.T) { testErr := errors.New("test") d := poolDialer{t: t} config1 := testConfig config1.Config = &pool.Config{ Active: 1, Idle: 1, Wait: true, } p := NewPool(config1) p.Slice.New = func(ctx context.Context) (io.Closer, error) { return d.dial() } defer p.Close() c := p.Get(context.TODO()) errs := startGoroutines(p, "ERR", testErr) d.check("before close", p, 1, 1) d.dialErr = errors.New("dial") c.Close() nilCount := 0 errCount := 0 timeout := time.After(2 * time.Second) for i := 0; i < cap(errs); i++ { select { case err := <-errs: switch err { case nil: nilCount++ case d.dialErr: errCount++ default: t.Fatalf("expected dial error or nil, got %v", err) } case <-timeout: t.Logf("Wait all the time and timeout %d", i) return } } if nilCount != 1 { t.Errorf("expected one nil error, got %d", nilCount) } if errCount != cap(errs)-1 { t.Errorf("expected %d dial erors, got %d", cap(errs)-1, errCount) } d.check("done", p, cap(errs), 0) } func BenchmarkPoolGet(b *testing.B) { b.StopTimer() p := NewPool(testConfig) c := p.Get(context.Background()) if err := c.Err(); err != nil { b.Fatal(err) } c.Close() defer p.Close() b.StartTimer() for i := 0; i < b.N; i++ { c := p.Get(context.Background()) c.Close() } } func BenchmarkPoolGetErr(b *testing.B) { b.StopTimer() p := NewPool(testConfig) c := p.Get(context.Background()) if err := c.Err(); err != nil { b.Fatal(err) } c.Close() defer p.Close() b.StartTimer() for i := 0; i < b.N; i++ { c = p.Get(context.Background()) if err := c.Err(); err != nil { b.Fatal(err) } c.Close() } } func BenchmarkPoolGetPing(b *testing.B) { b.StopTimer() p := NewPool(testConfig) c := p.Get(context.Background()) if err := c.Err(); err != nil { b.Fatal(err) } c.Close() defer p.Close() b.StartTimer() for i := 0; i < b.N; i++ { c := p.Get(context.Background()) if _, err := c.Do("PING"); err != nil { b.Fatal(err) } c.Close() } } func BenchmarkPooledConn(b *testing.B) { p := NewPool(testConfig) defer p.Close() for i := 0; i < b.N; i++ { ctx := context.TODO() c := p.Get(ctx) c2 := c.WithContext(context.TODO()) if _, err := c2.Do("PING"); err != nil { b.Fatal(err) } c2.Close() } }