Skip to content

Commit c1dfb99

Browse files
everpeaceeverpeace
authored
Merge pull request #26 from pfnet-research/get-task-in-chunks
get tasks in chunks in get-task command
2 parents bfb8681 + 6e9ae18 commit c1dfb99

File tree

5 files changed

+64
-42
lines changed

5 files changed

+64
-42
lines changed

‎cmd/root.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,9 @@ func init() {
205205
)
206206
viperBindPFlag("Redis.IdleCheckFrequency", cmdOpts.Redis.IdleCheckFrequency.String(), flag.Lookup("redis-idle-check-frequency"))
207207

208+
flag.Int("redis-chunk-size-in-get", cmdOpts.Redis.ChunkSizeInGet, "chunk size when gettings tasks from redis")
209+
viperBindPFlag("Redis.ChunkSizeInGet", strconv.Itoa(cmdOpts.Redis.ChunkSizeInGet), flag.Lookup("redis-chunk-size-in-get"))
210+
208211
// BackoffConfig for redis
209212
flag.Duration("redis-backoff-initial-interval", cmdOpts.Redis.Backoff.InitialInterval, "initial interval of exponential backoff used in redis operation")
210213
viperBindPFlag("Redis.Backoff.InitialInterval", cmdOpts.Redis.Backoff.InitialInterval.String(), flag.Lookup("redis-backoff-initial-interval"))
@@ -294,9 +297,10 @@ func mustInitializeQueueBackend() {
294297
queueBackend, err = backendfactory.NewBackend(logger, backendconfig.Config{
295298
BackendType: cmdOpts.Backend,
296299
Redis: &backendconfig.RedisConfig{
297-
KeyPrefix: cmdOpts.Redis.KeyPrefix,
298-
Client: cmdOpts.Redis.NewClient(),
299-
Backoff: cmdOpts.Redis.Backoff,
300+
KeyPrefix: cmdOpts.Redis.KeyPrefix,
301+
Client: cmdOpts.Redis.NewClient(),
302+
Backoff: cmdOpts.Redis.Backoff,
303+
ChunkSizeInGet: cmdOpts.Redis.ChunkSizeInGet,
300304
},
301305
})
302306

‎pkg/backend/config/config.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,10 @@ type Config struct {
3030
}
3131

3232
type RedisConfig struct {
33-
KeyPrefix string
34-
Client *redis.Client
35-
Backoff BackoffConfig
33+
KeyPrefix string
34+
Client *redis.Client
35+
Backoff BackoffConfig
36+
ChunkSizeInGet int
3637
}
3738

3839
// TODO: support UniversalOptions
@@ -41,16 +42,17 @@ type RedisClientConfig struct {
4142
Password string `json:"password" yaml:"password" default:""`
4243
DB int `json:"db" yaml:"db" default:"0"`
4344

44-
DialTimeout time.Duration `json:"dialTimeout" yaml:"dialTimeout" default:"30s"`
45-
ReadTimeout time.Duration `json:"readTimeout" yaml:"readTimeout" default:"10m"`
46-
WriteTimeout time.Duration `json:"writeTimeout" yaml:"writeTimeout" default:"10m"`
47-
45+
DialTimeout time.Duration `json:"dialTimeout" yaml:"dialTimeout" default:"30s"`
46+
ReadTimeout time.Duration `json:"readTimeout" yaml:"readTimeout" default:"10m"`
47+
WriteTimeout time.Duration `json:"writeTimeout" yaml:"writeTimeout" default:"10m"`
4848
PoolSize int `json:"poolSize" yaml:"poolSize" default:"-"`
4949
MinIdleConns int `json:"minIdleConns" yaml:"minIdleConns" default:"-"`
5050
MaxConnAge time.Duration `json:"maxConnAge" yaml:"maxConnAge" default:"-"`
5151
PoolTimeout time.Duration `json:"poolTimeout" yaml:"poolTimeout" default:"-"`
5252
IdleTimeout time.Duration `json:"idleTimeout" yaml:"idleTimeout" default:"5m"`
5353
IdleCheckFrequency time.Duration `json:"idleCheckFrequency" yaml:"idleCheckFrequency" default:"1m"`
54+
55+
ChunkSizeInGet int `json:"chunkSizeInGet" yaml:"chunkSizeInGet" default:"10000"`
5456
}
5557

5658
func (c RedisClientConfig) NewClient() *redis.Client {

‎pkg/backend/redis/redis_test.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,10 @@ var _ = Describe("Backend", func() {
106106
ibackend, err := NewBackend(logger, backendconfig.Config{
107107
BackendType: "redis",
108108
Redis: &backendconfig.RedisConfig{
109-
KeyPrefix: "test",
110-
Client: client,
111-
Backoff: backoffConfig,
109+
KeyPrefix: "test",
110+
Client: client,
111+
Backoff: backoffConfig,
112+
ChunkSizeInGet: 1000,
112113
},
113114
})
114115
Expect(err).NotTo(HaveOccurred())

‎pkg/backend/redis/task.go‎

Lines changed: 41 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -111,39 +111,53 @@ func (b *Backend) getTasksByUIDs(queueUID string, taskUIDs []string, filter func
111111
if len(taskUIDs) == 0 {
112112
return []*task.Task{}, nil
113113
}
114-
taskKeys := []string{}
115-
for _, uid := range taskUIDs {
116-
taskKeys = append(taskKeys, b.taskKey(queueUID, uid))
117-
}
118-
rawTasks, err := b.Client.MGet(taskKeys...).Result()
119-
if err != nil {
120-
return nil, err
114+
115+
chunks := [][]string{}
116+
chunkSize := b.RedisConfig.ChunkSizeInGet
117+
for i := 0; i < len(taskUIDs); i += chunkSize {
118+
until := len(taskUIDs)
119+
if i+chunkSize <= len(taskUIDs) {
120+
until = i + chunkSize
121+
}
122+
chunks = append(chunks, taskUIDs[i:until])
121123
}
124+
122125
tasks := []*task.Task{}
123-
for _, rawTask := range rawTasks {
124-
if rawTask == nil {
125-
lggr.Error().
126-
Interface("rawData", nil).
127-
Msg("Internal error. rawData is null. Skipped.")
128-
continue
126+
for _, chunk := range chunks {
127+
taskKeys := []string{}
128+
for _, uid := range chunk {
129+
taskKeys = append(taskKeys, b.taskKey(queueUID, uid))
129130
}
130-
var t task.Task
131-
rawTaskStr, ok := rawTask.(string)
132-
if !ok {
133-
lggr.Error().
134-
Interface("rawData", rawTask).
135-
Str("type", reflect.TypeOf(rawTask).String()).
136-
Msg("Internal error. rawData should be string. Skipped.")
137-
continue
138-
}
139-
if err := json.Unmarshal([]byte(rawTaskStr), &t); err != nil {
140-
lggr.Error().Str("data", rawTaskStr).Msg("Failed to unmarshal to Task. Skipped.")
141-
continue
131+
rawTasks, err := b.Client.MGet(taskKeys...).Result()
132+
if err != nil {
133+
return nil, err
142134
}
143-
if filter(&t) {
144-
tasks = append(tasks, &t)
135+
for _, rawTask := range rawTasks {
136+
if rawTask == nil {
137+
lggr.Error().
138+
Interface("rawData", nil).
139+
Msg("Internal error. rawData is null. Skipped.")
140+
continue
141+
}
142+
var t task.Task
143+
rawTaskStr, ok := rawTask.(string)
144+
if !ok {
145+
lggr.Error().
146+
Interface("rawData", rawTask).
147+
Str("type", reflect.TypeOf(rawTask).String()).
148+
Msg("Internal error. rawData should be string. Skipped.")
149+
continue
150+
}
151+
if err := json.Unmarshal([]byte(rawTaskStr), &t); err != nil {
152+
lggr.Error().Str("data", rawTaskStr).Msg("Failed to unmarshal to Task. Skipped.")
153+
continue
154+
}
155+
if filter(&t) {
156+
tasks = append(tasks, &t)
157+
}
145158
}
146159
}
160+
147161
return tasks, nil
148162
}
149163

‎pkg/worker/worker_test.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,9 @@ var _ = Describe("Worker", func() {
7474
bcknd, err = backendfactory.NewBackend(logger, backendconfig.Config{
7575
BackendType: "redis",
7676
Redis: &backendconfig.RedisConfig{
77-
Client: client,
78-
Backoff: backendConfig,
77+
Client: client,
78+
Backoff: backendConfig,
79+
ChunkSizeInGet: 1000,
7980
},
8081
})
8182
Expect(err).NotTo(HaveOccurred())

0 commit comments

Comments
 (0)