go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/featureBreaker/tq.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package featureBreaker
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	tq "go.chromium.org/luci/gae/service/taskqueue"
    22  )
    23  
    24  type tqState struct {
    25  	*state
    26  
    27  	c  context.Context
    28  	tq tq.RawInterface
    29  }
    30  
    31  var _ tq.RawInterface = (*tqState)(nil)
    32  
    33  func (t *tqState) AddMulti(tasks []*tq.Task, queueName string, cb tq.RawTaskCB) error {
    34  	if len(tasks) == 0 {
    35  		return nil
    36  	}
    37  	return t.run(t.c, func() (err error) { return t.tq.AddMulti(tasks, queueName, cb) })
    38  }
    39  
    40  func (t *tqState) DeleteMulti(tasks []*tq.Task, queueName string, cb tq.RawCB) error {
    41  	if len(tasks) == 0 {
    42  		return nil
    43  	}
    44  	return t.run(t.c, func() error { return t.tq.DeleteMulti(tasks, queueName, cb) })
    45  }
    46  
    47  func (t *tqState) Lease(maxTasks int, queueName string, leaseTime time.Duration) (tasks []*tq.Task, err error) {
    48  	err = t.run(t.c, func() (err error) {
    49  		tasks, err = t.tq.Lease(maxTasks, queueName, leaseTime)
    50  		return
    51  	})
    52  	if err != nil {
    53  		tasks = nil
    54  	}
    55  	return
    56  }
    57  
    58  func (t *tqState) LeaseByTag(maxTasks int, queueName string, leaseTime time.Duration, tag string) (tasks []*tq.Task, err error) {
    59  	err = t.run(t.c, func() (err error) {
    60  		tasks, err = t.tq.LeaseByTag(maxTasks, queueName, leaseTime, tag)
    61  		return
    62  	})
    63  	if err != nil {
    64  		tasks = nil
    65  	}
    66  	return
    67  }
    68  
    69  func (t *tqState) ModifyLease(task *tq.Task, queueName string, leaseTime time.Duration) error {
    70  	return t.run(t.c, func() error { return t.tq.ModifyLease(task, queueName, leaseTime) })
    71  }
    72  
    73  func (t *tqState) Purge(queueName string) error {
    74  	return t.run(t.c, func() error { return t.tq.Purge(queueName) })
    75  }
    76  
    77  func (t *tqState) Stats(queueNames []string, cb tq.RawStatsCB) error {
    78  	return t.run(t.c, func() error { return t.tq.Stats(queueNames, cb) })
    79  }
    80  
    81  func (t *tqState) Constraints() tq.Constraints {
    82  	return t.tq.Constraints()
    83  }
    84  
    85  func (t *tqState) GetTestable() tq.Testable {
    86  	return t.tq.GetTestable()
    87  }
    88  
    89  // FilterTQ installs a featureBreaker TaskQueue filter in the context.
    90  func FilterTQ(c context.Context, defaultError error) (context.Context, FeatureBreaker) {
    91  	state := newState(defaultError)
    92  	return tq.AddRawFilters(c, func(ic context.Context, tq tq.RawInterface) tq.RawInterface {
    93  		return &tqState{state, ic, tq}
    94  	}), state
    95  }