bitleak / go-redis-pool Goto Github PK

// NewHAConnFactory create new ha factory
func newClientPool(cfg *HAConfig) *clientPool {
	slavePassword := cfg.Password
	if cfg.ReadonlyPassword != "" {
		slavePassword = cfg.ReadonlyPassword
	}

	pool := &clientPool{
		pollType:           cfg.PollType,
		autoEjectHost:      cfg.AutoEjectHost,
		serverFailureLimit: cfg.ServerFailureLimit,
		serverRetryTimeout: cfg.ServerRetryTimeout,
		minServerNum:       cfg.MinServerNum,

		stopCh: make(chan struct{}, 0),
		rand:   rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	pool.slaves = make([]*client, len(cfg.Slaves))
	options := cfg.Options
	for i, slave := range cfg.Slaves {
		slaveOptions := *options
		slaveOptions.Addr = slave  //  Do you need to handle the extra colons for PollByWeight here
		slaveOptions.Password = slavePassword
		redisCli := redis.NewClient(&slaveOptions)
		cli := newClient(redisCli, cfg.weights[i])
		if cfg.AutoEjectHost && len(cfg.Slaves) > 1 {
			redisCli.AddHook(newFailureHook(cli))
		}
		pool.slaves[i] = cli
	}
	pool.alives = pool.slaves
	if pool.pollType == PollByWeight {
		weightRanges := make([]int64, len(pool.alives))
		weightRanges[0] = pool.alives[0].weight
		for i := 1; i < len(pool.alives); i++ {
			weightRanges[i] = weightRanges[i-1] + pool.alives[i].weight
		}
		pool.weightRanges = weightRanges
	}
	if cfg.AutoEjectHost && len(cfg.Slaves) > 1 {
		go pool.detectFailureTick()
	}
	return pool
}

// MGetWithGD is like MGet but returns all the values that it managed to get
func (p *Pool) MGetWithGD(ctx context.Context, keys ...string) ([]interface{}, map[string]error) {
	keyErrors := make(map[string]error, 0)

	if _, ok := p.connFactory.(*HAConnFactory); ok {
		//conn, _ := p.connFactory.getMasterConn()
		conn, _ := p.connFactory.getSlaveConn(keys...)
		vals, err := conn.MGet(ctx, keys...).Result()
		if err != nil {
			for _, key := range keys {
				keyErrors[key] = err
			}
			return nil, keyErrors
		}
		return vals, nil
	}

	fn := func(factory *ShardConnFactory, keyList ...string) redis.Cmder {
		conn, err := factory.getSlaveConn(keyList[0])
		if err != nil {
			args := make([]interface{}, 1+len(keyList))
			args[0] = "mget"
			for i, key := range keyList {
				args[1+i] = key
			}
			cmd := redis.NewSliceCmd(ctx, args...)
			cmd.SetErr(err)
			return cmd
		}
		return conn.MGet(ctx, keyList...)
	}

	factory := p.connFactory.(*ShardConnFactory)
	results := factory.doMultiKeys(fn, keys...)
	keyVals := make(map[string]interface{}, 0)
	for _, result := range results {
		args := result.Args()
		vals, err := result.(*redis.SliceCmd).Result()
		if err != nil {
			for i, arg := range args {
				if i == 0 {
					continue
				}
				keyErrors[arg.(string)] = err
			}
			continue
		}
		for i, val := range vals {
			keyVals[args[i+1].(string)] = val
		}
	}
	vals := make([]interface{}, len(keys))
	for i, key := range keys {
		vals[i] = nil
		if val, ok := keyVals[key]; ok {
			vals[i] = val
		}
	}
	return vals, keyErrors
}