OLD | NEW |
1 // Copyright 2015 The Chromium Authors. All rights reserved. | 1 // Copyright 2015 The Chromium Authors. All rights reserved. |
2 // Use of this source code is governed by a BSD-style license that can be | 2 // Use of this source code is governed by a BSD-style license that can be |
3 // found in the LICENSE file. | 3 // found in the LICENSE file. |
4 | 4 |
5 package taskqueue | 5 package taskqueue |
6 | 6 |
7 import ( | 7 import ( |
8 "golang.org/x/net/context" | 8 "golang.org/x/net/context" |
9 ) | 9 ) |
10 | 10 |
11 type key int | 11 type key int |
12 | 12 |
13 var ( | 13 var ( |
14 taskQueueKey key | 14 taskQueueKey key |
15 taskQueueFilterKey key = 1 | 15 taskQueueFilterKey key = 1 |
16 ) | 16 ) |
17 | 17 |
18 // RawFactory is the function signature for RawFactory methods compatible with | 18 // RawFactory is the function signature for RawFactory methods compatible with |
19 // SetRawFactory. | 19 // SetRawFactory. wantTxn is true if the Factory should return the taskqueue in |
20 type RawFactory func(context.Context) RawInterface | 20 // the current transaction, and false if the Factory should return the |
| 21 // non-transactional (root) taskqueue service. |
| 22 type RawFactory func(c context.Context, wantTxn bool) RawInterface |
21 | 23 |
22 // RawFilter is the function signature for a RawFilter TQ implementation. It | 24 // RawFilter is the function signature for a RawFilter TQ implementation. It |
23 // gets the current TQ implementation, and returns a new TQ implementation | 25 // gets the current TQ implementation, and returns a new TQ implementation |
24 // backed by the one passed in. | 26 // backed by the one passed in. |
25 type RawFilter func(context.Context, RawInterface) RawInterface | 27 type RawFilter func(context.Context, RawInterface) RawInterface |
26 | 28 |
27 // getUnfiltered gets gets the RawInterface implementation from context without | 29 // getUnfiltered gets gets the RawInterface implementation from context without |
28 // any of the filters applied. | 30 // any of the filters applied. |
29 func getUnfiltered(c context.Context) RawInterface { | 31 func getUnfiltered(c context.Context, wantTxn bool) RawInterface { |
30 if f, ok := c.Value(taskQueueKey).(RawFactory); ok && f != nil { | 32 if f, ok := c.Value(taskQueueKey).(RawFactory); ok && f != nil { |
31 » » return f(c) | 33 » » return f(c, wantTxn) |
32 } | 34 } |
33 return nil | 35 return nil |
34 } | 36 } |
35 | 37 |
36 // GetRaw gets the RawInterface implementation from context. | 38 // getFiltered gets the taskqueue (transactional or not), and applies all of |
37 func GetRaw(c context.Context) RawInterface { | 39 // the currently installed filters to it. |
38 » ret := getUnfiltered(c) | 40 func getFiltered(c context.Context, wantTxn bool) RawInterface { |
| 41 » ret := getUnfiltered(c, wantTxn) |
39 if ret == nil { | 42 if ret == nil { |
40 return nil | 43 return nil |
41 } | 44 } |
42 for _, f := range getCurFilters(c) { | 45 for _, f := range getCurFilters(c) { |
43 ret = f(c, ret) | 46 ret = f(c, ret) |
44 } | 47 } |
45 return ret | 48 return ret |
46 } | 49 } |
47 | 50 |
| 51 // GetRaw gets the RawInterface implementation from context. |
| 52 func GetRaw(c context.Context) RawInterface { |
| 53 return getFiltered(c, true) |
| 54 } |
| 55 |
| 56 // GetRawNoTxn gets the RawInterface implementation from context. If there's a |
| 57 // currently active transaction, this will return a non-transactional connection |
| 58 // to the taskqueue, otherwise this is the same as GetRaw. |
| 59 func GetRawNoTxn(c context.Context) RawInterface { |
| 60 return getFiltered(c, false) |
| 61 } |
| 62 |
| 63 // Get gets the Interface implementation from context. |
| 64 func Get(c context.Context) Interface { |
| 65 return &taskqueueImpl{GetRaw(c)} |
| 66 } |
| 67 |
| 68 // GetNoTxn gets the Interface implementation from context. |
| 69 func GetNoTxn(c context.Context) Interface { |
| 70 return &taskqueueImpl{GetRawNoTxn(c)} |
| 71 } |
| 72 |
48 // SetRawFactory sets the function to produce RawInterface instances, as returne
d by | 73 // SetRawFactory sets the function to produce RawInterface instances, as returne
d by |
49 // the GetRaw method. | 74 // the GetRaw method. |
50 func SetRawFactory(c context.Context, tqf RawFactory) context.Context { | 75 func SetRawFactory(c context.Context, tqf RawFactory) context.Context { |
51 return context.WithValue(c, taskQueueKey, tqf) | 76 return context.WithValue(c, taskQueueKey, tqf) |
52 } | 77 } |
53 | 78 |
54 // SetRaw sets the current RawInterface object in the context. Useful for testin
g | 79 // SetRaw sets the current RawInterface object in the context. Useful for testin
g |
55 // with a quick mock. This is just a shorthand SetRawFactory invocation to SetRa
w | 80 // with a quick mock. This is just a shorthand SetRawFactory invocation to SetRa
w |
56 // a RawFactory which always returns the same object. | 81 // a RawFactory which always returns the same object. |
57 func SetRaw(c context.Context, tq RawInterface) context.Context { | 82 func SetRaw(c context.Context, tq RawInterface) context.Context { |
58 » return SetRawFactory(c, func(context.Context) RawInterface { return tq }
) | 83 » return SetRawFactory(c, func(context.Context, bool) RawInterface { retur
n tq }) |
59 } | 84 } |
60 | 85 |
61 func getCurFilters(c context.Context) []RawFilter { | 86 func getCurFilters(c context.Context) []RawFilter { |
62 curFiltsI := c.Value(taskQueueFilterKey) | 87 curFiltsI := c.Value(taskQueueFilterKey) |
63 if curFiltsI != nil { | 88 if curFiltsI != nil { |
64 return curFiltsI.([]RawFilter) | 89 return curFiltsI.([]RawFilter) |
65 } | 90 } |
66 return nil | 91 return nil |
67 } | 92 } |
68 | 93 |
69 // AddRawFilters adds RawInterface filters to the context. | 94 // AddRawFilters adds RawInterface filters to the context. |
70 func AddRawFilters(c context.Context, filts ...RawFilter) context.Context { | 95 func AddRawFilters(c context.Context, filts ...RawFilter) context.Context { |
71 if len(filts) == 0 { | 96 if len(filts) == 0 { |
72 return c | 97 return c |
73 } | 98 } |
74 cur := getCurFilters(c) | 99 cur := getCurFilters(c) |
75 newFilts := make([]RawFilter, 0, len(cur)+len(filts)) | 100 newFilts := make([]RawFilter, 0, len(cur)+len(filts)) |
76 newFilts = append(newFilts, getCurFilters(c)...) | 101 newFilts = append(newFilts, getCurFilters(c)...) |
77 newFilts = append(newFilts, filts...) | 102 newFilts = append(newFilts, filts...) |
78 return context.WithValue(c, taskQueueFilterKey, newFilts) | 103 return context.WithValue(c, taskQueueFilterKey, newFilts) |
79 } | 104 } |
OLD | NEW |