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