| 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 |