├── .gitignore ├── Dockerfile ├── README.md ├── client.go ├── go.mod ├── go.sum ├── task.go ├── tasq-client-python ├── README.md ├── setup.py └── tasq_client │ ├── __init__.py │ ├── check_type.py │ ├── check_type_test.py │ └── client.py ├── tasq-rate-estimate └── main.go ├── tasq-server ├── homepage.go ├── json.go ├── main.go ├── queues.go ├── rate_tracker.go ├── rate_tracker_test.go └── task.go └── tasq-transfer └── main.go /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.egg-info/ 3 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM golang:1.18-alpine 2 | 3 | WORKDIR /app 4 | COPY go.mod ./ 5 | COPY go.sum ./ 6 | COPY tasq-server/*.go ./ 7 | RUN go mod download 8 | 9 | EXPOSE 8080 10 | ENTRYPOINT ["go", "run", "."] 11 | CMD ["-addr=:8080"] -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # tasq 2 | 3 | **tasq** is a simple HTTP-based task queue. Each task is represented as a string (it could be anything). 4 | 5 | Tasks are pushed to the queue via an HTTP endpoint, and popped using another endpoint. The worker performing the popped task can signal its completion using a third endpoint. If tasks aren't completed after a timeout, they can be popped again once all other pending tasks have been popped. 6 | 7 | # Protocol 8 | 9 | Here are endpoints for pushing and popping tasks: 10 | 11 | * `/task/push` - add a task to the queue. Simply provide a `?contents=X` query argument. 12 | * `/task/push_batch` - POST to this endpoint with a JSON array of tasks. For example, `["hi", "test"]`. 13 | * `/task/pop` - pop a task from the queue. If no tasks are available, this may indicate a timeout after which the longest-running task would timeout. 14 | * On normal response, will return something like `{"data": {"id": "...", "contents": "..."}}`. 15 | * If queue is empty, will return something like `{"data": {"done": false, "retry": 3.14}}`, where `retry` is the number of seconds after which to try popping again, and `done` is `true` if no tasks are pending or running. 16 | * `/task/completed` - indicate that the task is completed. Simply provide a `?id=X` query argument. 17 | 18 | Additionally, these are some endpoints that may be helpful for maintaining a running queue in practice: 19 | * `/` - an overview of all the queues, with some buttons and forms to quickly manipulate queues. 20 | * `/summary` - a textual overview of all the queues. 21 | * `/counts` - get a dictionary containing sizes of queues. Has keys `pending`, `running`, `expired`, and `completed`. 22 | * `/task/peek` - look at the next task that would be returned by `/task/pop`. When the queue is empty but tasks are still in progress (but not timed out), this returns extra information. In addition to `done` and `retry` fields, this will return a `next` field containing a dictionary with `id` and `contents` of the next task that will expire. This can make it easier for a human to see which tasks are repeatedly failing or timing out. 23 | * `/task/clear` - delete all pending and running tasks in the queue. 24 | * `/task/expire_all` - set all currently running tasks as expired so that they can be re-popped immediately. 25 | * `/task/queue_expired` - move all expired tasks from the `in-progress` queue to the `pending` queue. This used to be helpful when the `/counts` endpoint didn't count expired tasks, but it will also have an effect on prematurely expired tasks: if any worker was still working on an expired task and calls `/task/completed`, a task in the `pending` queue will not be successfully marked as completed. 26 | 27 | # Persistence 28 | 29 | Using the `-save-path` and `-save-interval` flags, you can configure `tasq-server` to periodically dump its state to a file. This can prevent long-running jobs from losing progress if the server crashes or restarts. 30 | 31 | When using file persistence, it is possible that some progress will be lost when the server restarts. If tasks were pushed between the latest save and the restart, then these tasks will be lost. If tasks were completed during this interval, then the tasks will reappear in the queue upon restart. To solve the latter issue, one can make workers able to handle already-completed tasks. Solving the former issue is more difficult in general, but it is unlikely to be a problem for jobs where all work is queued at the start and then gradually worked through by workers. 32 | -------------------------------------------------------------------------------- /client.go: -------------------------------------------------------------------------------- 1 | package tasq 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | "io" 7 | "net/http" 8 | "net/url" 9 | "path" 10 | "strconv" 11 | "strings" 12 | "time" 13 | 14 | "github.com/pkg/errors" 15 | ) 16 | 17 | const DefaultKeepaliveInterval = time.Second * 30 18 | 19 | // A Task stores information about a popped task. 20 | type Task struct { 21 | ID string `json:"id"` 22 | Contents string `json:"contents"` 23 | } 24 | 25 | // QueueCounts stores the number of in-progress, pending, and completed tasks. 26 | type QueueCounts struct { 27 | Completed int64 `json:"completed"` 28 | Pending int64 `json:"pending"` 29 | Expired int64 `json:"expired"` 30 | Running int64 `json:"running"` 31 | } 32 | 33 | // A Client makes API calls to a tasq server. 34 | // 35 | // The server is identified as a URL. For example, you might provide a parsed 36 | // URL "http://myserver.com:8080". The path in the URL is appended with API 37 | // endpoint paths, but the protocol, host, and port are retained. 38 | type Client struct { 39 | URL *url.URL 40 | 41 | // Provide to enable basic auth. 42 | Username string 43 | Password string 44 | 45 | // KeepaliveInterval is used for the keepalive Goroutine created by the 46 | // PopRunningTask method. Defaults to DefaultKeepaliveInterval. 47 | KeepaliveInterval time.Duration 48 | } 49 | 50 | // NewClient creates a client with a base server URL. 51 | // 52 | // Optionally, a context name can be passed to scope the task queue, 53 | // as well as a username and password. 54 | // 55 | // Returns an error if the URL fails to parse. 56 | func NewClient(baseURL string, contextUserPass ...string) (*Client, error) { 57 | if len(contextUserPass) != 1 && len(contextUserPass) != 3 { 58 | panic("zero or one context arguments expected") 59 | } 60 | parsed, err := url.Parse(baseURL) 61 | if err != nil { 62 | return nil, errors.Wrap(err, "new client") 63 | } 64 | if len(contextUserPass) > 0 { 65 | parsed.RawQuery = (url.Values{"context": contextUserPass[:1]}).Encode() 66 | } 67 | res := &Client{URL: parsed} 68 | if len(contextUserPass) == 3 { 69 | res.Username = contextUserPass[1] 70 | res.Password = contextUserPass[2] 71 | } 72 | return res, nil 73 | } 74 | 75 | // Push adds a task to the queue and returns its ID. 76 | func (c *Client) Push(contents string) (string, error) { 77 | var response string 78 | err := c.postForm("/task/push", "contents", contents, &response) 79 | return response, err 80 | } 81 | 82 | // PushBatch adds a batch of tasks to the queue and return their IDs. 83 | func (c *Client) PushBatch(contents []string) ([]string, error) { 84 | var response []string 85 | err := c.postJSON("/task/push_batch", contents, &response) 86 | return response, err 87 | } 88 | 89 | // Pop retrieves a pending task from the queue. 90 | // 91 | // If no task is returned, a retry time may be returned indicating the number 92 | // of seconds until the next in-progress task will expire. If this retry time 93 | // is also nil, then the queue has been exhausted. 94 | func (c *Client) Pop() (*Task, *float64, error) { 95 | var response struct { 96 | ID *string `json:"id"` 97 | Contents *string `json:"contents"` 98 | Done bool `json:"done"` 99 | Retry float64 `json:"retry"` 100 | } 101 | if err := c.get("/task/pop", &response); err != nil { 102 | return nil, nil, err 103 | } 104 | if response.ID != nil && response.Contents != nil { 105 | return &Task{ID: *response.ID, Contents: *response.Contents}, nil, nil 106 | } else if response.Done { 107 | return nil, nil, nil 108 | } else { 109 | return nil, &response.Retry, nil 110 | } 111 | } 112 | 113 | // PopBatch retrieves at most n tasks from the queue. 114 | // 115 | // If fewer than n tasks are returned, then a retry time (in seconds) may be 116 | // returned to indicate when the next pending task will expire. 117 | // 118 | // If no tasks are returned and the retry time is nil, then the queue has been 119 | // exhausted. 120 | func (c *Client) PopBatch(n int) ([]*Task, *float64, error) { 121 | var response struct { 122 | Done bool `json:"done"` 123 | Retry float64 `json:"retry"` 124 | Tasks []*Task `json:"tasks"` 125 | } 126 | if err := c.postForm("/task/pop_batch", "count", strconv.Itoa(n), &response); err != nil { 127 | return nil, nil, err 128 | } 129 | if response.Done { 130 | return nil, nil, nil 131 | } else { 132 | return response.Tasks, &response.Retry, nil 133 | } 134 | } 135 | 136 | // PopRunningTask pops a task from the queue, potentially blocking until a task 137 | // becomes available, and returns a new *RunningTask. 138 | // 139 | // If no tasks are pending, nil is returned. 140 | // 141 | // If a *RunningTask is successfully returned, the caller must call Completed() 142 | // or Cancel() on it to clean up resources. 143 | func (c *Client) PopRunningTask() (*RunningTask, error) { 144 | for { 145 | task, wait, err := c.Pop() 146 | if err != nil { 147 | return nil, err 148 | } else if task != nil { 149 | interval := c.KeepaliveInterval 150 | if interval == 0 { 151 | interval = DefaultKeepaliveInterval 152 | } 153 | return newRunningTask(c, task.Contents, task.ID, interval), nil 154 | } else if wait != nil { 155 | time.Sleep(time.Duration(float64(time.Second) * (*wait))) 156 | } else { 157 | return nil, nil 158 | } 159 | } 160 | } 161 | 162 | // Completed tells the server that the identified task was completed. 163 | func (c *Client) Completed(id string) error { 164 | return c.postForm("/task/completed", "id", id, nil) 165 | } 166 | 167 | // CompletedBatch tells the server that the identified tasks were completed. 168 | func (c *Client) CompletedBatch(ids []string) error { 169 | return c.postJSON("/task/completed_batch", ids, nil) 170 | } 171 | 172 | // Completed tells the server to restart the timeout window for an in-progress 173 | // task. 174 | func (c *Client) Keepalive(id string) error { 175 | return c.postForm("/task/keepalive", "id", id, nil) 176 | } 177 | 178 | // QueueCounts gets the number of tasks in each queue. 179 | func (c *Client) QueueCounts() (*QueueCounts, error) { 180 | var result QueueCounts 181 | if err := c.get("/counts", &result); err != nil { 182 | return nil, err 183 | } 184 | return &result, nil 185 | } 186 | 187 | func (c *Client) get(path string, output interface{}) error { 188 | reqURL := c.urlForPath(path) 189 | req, err := http.NewRequest("GET", reqURL.String(), nil) 190 | if err != nil { 191 | return errors.Wrap(err, "get "+path) 192 | } 193 | if c.Username != "" || c.Password != "" { 194 | req.SetBasicAuth(c.Username, c.Password) 195 | } 196 | resp, err := http.DefaultClient.Do(req) 197 | if err := c.handleResponse(resp, err, output); err != nil { 198 | return errors.Wrap(err, "get "+path) 199 | } 200 | return nil 201 | } 202 | 203 | func (c *Client) postForm(path, key, value string, output interface{}) error { 204 | postBody := strings.NewReader(url.QueryEscape(key) + "=" + url.QueryEscape(value)) 205 | return c.post(path, "application/x-www-form-urlencoded", postBody, output) 206 | } 207 | 208 | func (c *Client) postJSON(path string, input, output interface{}) error { 209 | data, err := json.Marshal(input) 210 | if err != nil { 211 | return errors.Wrap(err, "post "+path) 212 | } 213 | return c.post(path, "application/json", bytes.NewReader(data), output) 214 | } 215 | 216 | func (c *Client) post(path string, contentType string, input io.Reader, output interface{}) error { 217 | reqURL := c.urlForPath(path) 218 | req, err := http.NewRequest("POST", reqURL.String(), input) 219 | if err != nil { 220 | return errors.Wrap(err, "get "+path) 221 | } 222 | req.Header.Set("content-type", contentType) 223 | if c.Username != "" || c.Password != "" { 224 | req.SetBasicAuth(c.Username, c.Password) 225 | } 226 | resp, err := http.DefaultClient.Do(req) 227 | if err := c.handleResponse(resp, err, output); err != nil { 228 | return errors.Wrap(err, "post "+path) 229 | } 230 | return nil 231 | } 232 | 233 | func (c *Client) handleResponse(resp *http.Response, err error, output interface{}) error { 234 | if err != nil { 235 | return err 236 | } 237 | defer resp.Body.Close() 238 | 239 | var response struct { 240 | Error *string `json:"error"` 241 | Data interface{} `json:"data"` 242 | } 243 | response.Data = output 244 | if err := json.NewDecoder(resp.Body).Decode(&response); err != nil { 245 | return err 246 | } else if response.Error != nil { 247 | return errors.New("remote error: " + *response.Error) 248 | } else { 249 | return nil 250 | } 251 | } 252 | 253 | func (c *Client) urlForPath(p string) *url.URL { 254 | u := *c.URL 255 | if u.Path == "/" || u.Path == "" { 256 | u.Path = p 257 | } else { 258 | u.Path = path.Join(u.Path, p) 259 | } 260 | return &u 261 | } 262 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/unixpickle/tasq 2 | 3 | go 1.14 4 | 5 | require ( 6 | github.com/pkg/errors v0.9.1 7 | github.com/unixpickle/essentials v1.3.0 8 | ) 9 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= 2 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 3 | github.com/unixpickle/essentials v1.3.0 h1:H258Z5Uo1pVzFjxD2rwFWzHPN3s0J0jLs5kuxTRSfCs= 4 | github.com/unixpickle/essentials v1.3.0/go.mod h1:dQ1idvqrgrDgub3mfckQm7osVPzT3u9rB6NK/LEhmtQ= 5 | -------------------------------------------------------------------------------- /task.go: -------------------------------------------------------------------------------- 1 | package tasq 2 | 3 | import ( 4 | "sync" 5 | "time" 6 | ) 7 | 8 | // A RunningTask represents an in-progress task that is actively being 9 | // performed by this process. Call Completed() on the object to mark it as 10 | // complete. 11 | // 12 | // The object will automatically manage a background Goroutine that sends 13 | // keepalives to the server until Completed() or Cancel() is called on it. 14 | type RunningTask struct { 15 | Contents string 16 | ID string 17 | 18 | client *Client 19 | 20 | cancelLock sync.Mutex 21 | cancelled bool 22 | cancelChan chan struct{} 23 | } 24 | 25 | func newRunningTask(client *Client, contents, id string, interval time.Duration) *RunningTask { 26 | r := &RunningTask{ 27 | Contents: contents, 28 | ID: id, 29 | client: client, 30 | cancelChan: make(chan struct{}), 31 | } 32 | go r.keepaliveLoop(interval) 33 | return r 34 | } 35 | 36 | // Completed marks the task as complete and cancels the keepalive loop. 37 | // 38 | // Even if this returns an error, the keepalive loop will be stopped. 39 | func (r *RunningTask) Completed() error { 40 | r.Cancel() 41 | return r.client.Completed(r.ID) 42 | } 43 | 44 | // Cancel the task's keepalive loop. 45 | // 46 | // This may be called any number of times, even if the task was completed, 47 | // in which case it will have no effect after the first cancellation. 48 | func (r *RunningTask) Cancel() { 49 | r.cancelLock.Lock() 50 | defer r.cancelLock.Unlock() 51 | if !r.cancelled { 52 | r.cancelled = true 53 | close(r.cancelChan) 54 | } 55 | } 56 | 57 | func (r *RunningTask) keepaliveLoop(interval time.Duration) { 58 | for { 59 | select { 60 | case <-time.After(interval): 61 | case <-r.cancelChan: 62 | return 63 | } 64 | r.client.Keepalive(r.ID) 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /tasq-client-python/README.md: -------------------------------------------------------------------------------- 1 | ## Tasq python client 2 | 3 | ### Initialization: 4 | 5 | ```python 6 | from tasq_client import TasqClient 7 | TasqClient(base_url="", 8 | username="", 9 | password="", 10 | context="") 11 | ``` 12 | 13 | ### Sample enqueue usage: 14 | 15 | ```python 16 | client.push(json.dumps({'key': 'value'})) 17 | ``` 18 | 19 | ### Sample worker: 20 | 21 | ```python 22 | while True: 23 | with client.pop_running_task() as task: 24 | if task is None: 25 | break 26 | d = json.loads(task.contents) 27 | assert d['key'] == 'value' 28 | ``` 29 | 30 | ## Troubleshooting 31 | 32 | If you get the following error: 33 | ```requests.exceptions.SSLError: HTTPSConnectionPool(host='', port=443): Max retries exceeded with url: /task/completed?context=test (Caused by SSLError(SSLError(1, '[SSL: DECRYPTION_FAILED_OR_BAD_RECORD_MAC] decryption failed or bad record mac (_ssl.c:2633)')))``` 34 | 35 | Try setting the following option at the start of your program: 36 | ```python 37 | import multiprocessing 38 | multiprocessing.set_start_method('spawn') 39 | ``` 40 | -------------------------------------------------------------------------------- /tasq-client-python/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="tasq-client-python", 5 | packages=["tasq_client"], 6 | install_requires=["requests"], 7 | version="0.1.17", 8 | ) 9 | -------------------------------------------------------------------------------- /tasq-client-python/tasq_client/__init__.py: -------------------------------------------------------------------------------- 1 | from .client import ( 2 | QueueCounts, 3 | RunningTask, 4 | Task, 5 | TasqClient, 6 | TasqMisbehavingServerError, 7 | TasqRemoteError, 8 | ) 9 | 10 | __all__ = [ 11 | "QueueCounts", 12 | "RunningTask", 13 | "TasqClient", 14 | "Task", 15 | "TasqMisbehavingServerError", 16 | "TasqRemoteError", 17 | ] 18 | -------------------------------------------------------------------------------- /tasq-client-python/tasq_client/check_type.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Any, Callable 3 | 4 | 5 | @dataclass(frozen=True, eq=True) 6 | class OptionalKey: 7 | """ 8 | An object to use as a key in a type template to indicate that the field 9 | need not be present in a dict. 10 | """ 11 | 12 | key: str 13 | 14 | 15 | @dataclass 16 | class OptionalValue: 17 | template: Any 18 | 19 | 20 | class CheckTypeException(Exception): 21 | """ 22 | An error indicating that the type of an object does not match the expected 23 | template for the object's type. 24 | """ 25 | 26 | 27 | def check_type(template: Any, obj: Any): 28 | """ 29 | Raise an exception if a given object does not match a type template. 30 | 31 | For example, a type template might be `str`, and `"hello"` would match the 32 | template whereas `1234` would not. 33 | 34 | Templates can include nested data structures. If a template is a dict, each 35 | key will map to a corresponding template for the value of that key. A key 36 | in a template must be present in the checked object, unless the key is 37 | wrapped in OptionalKey. 38 | 39 | If a template is a list, it should contain one object--the template for the 40 | elements of the list. 41 | 42 | The float template will accept both int and float types. 43 | """ 44 | # pylint: disable=cell-var-from-loop 45 | if isinstance(template, dict): 46 | if not isinstance(obj, dict): 47 | raise CheckTypeException(f"expected dict but got {type(obj)}") 48 | for k, v in template.items(): 49 | if isinstance(k, OptionalKey): 50 | k = k.key 51 | if k not in obj: 52 | continue 53 | if k not in obj: 54 | raise CheckTypeException(f"missing key: {k}") 55 | _wrap_check(f"value for key {k}", lambda: check_type(v, obj[k])) 56 | elif isinstance(template, list): 57 | if not isinstance(obj, list): 58 | raise CheckTypeException(f"expected list but got {type(obj)}") 59 | assert len(template) == 1 60 | value_template = template[0] 61 | for i, value in enumerate(obj): 62 | _wrap_check(f"value at index {i}", lambda: check_type(value_template, value)) 63 | elif template is float: 64 | if not isinstance(obj, int) and not isinstance(obj, float): 65 | raise CheckTypeException( 66 | f"expected type {template} to be float or int but got {type(obj)}" 67 | ) 68 | elif isinstance(template, OptionalValue): 69 | if obj is not None: 70 | _wrap_check("optional value", lambda: check_type(template.template, obj)) 71 | else: 72 | if not isinstance(obj, template): 73 | raise CheckTypeException(f"expected type {template} but got {type(obj)}") 74 | 75 | 76 | def _wrap_check(context: str, check_fn: Callable): 77 | try: 78 | check_fn() 79 | except CheckTypeException as exc: 80 | # pylint: disable=raise-missing-from 81 | raise CheckTypeException(f"{context}: {str(exc)}") 82 | -------------------------------------------------------------------------------- /tasq-client-python/tasq_client/check_type_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from .check_type import CheckTypeException, OptionalKey, check_type 4 | 5 | 6 | @pytest.mark.parametrize( 7 | ["template", "value", "expected"], 8 | [ 9 | (int, 3, True), 10 | (int, 3.0, False), 11 | (int, "hi", False), 12 | (float, 3, True), 13 | (float, 3.0, True), 14 | (float, "hi", False), 15 | ([str], [], True), 16 | ([str], ["hello"], True), 17 | ([str], ["hello", 3], False), 18 | ([str], 3, False), 19 | ([str], "hi", False), 20 | (dict(field=int), dict(), False), 21 | (dict(field=int), dict(field=3), True), 22 | (dict(field=int), dict(field=3, other=3), True), 23 | ({"field": int, OptionalKey("other"): str}, dict(field=3), True), 24 | ({"field": int, OptionalKey("other"): str}, dict(field=3, other="hi"), True), 25 | ({"field": int, OptionalKey("other"): str}, dict(field=3, other=3), False), 26 | (dict(field=dict(baz=int)), dict(field=3), False), 27 | (dict(field=dict(baz=int)), dict(field=dict(baz="hi")), False), 28 | (dict(field=dict(baz=int)), dict(field=dict(baz=3)), True), 29 | ], 30 | ) 31 | def test_check_type(template, value, expected): 32 | if expected: 33 | check_type(template, value) 34 | else: 35 | try: 36 | check_type(template, value) 37 | assert False, f"template {template} should not match {value}" 38 | except CheckTypeException: 39 | pass 40 | 41 | # Everything should be an object. 42 | check_type(object, value) 43 | -------------------------------------------------------------------------------- /tasq-client-python/tasq_client/client.py: -------------------------------------------------------------------------------- 1 | import random 2 | import sys 3 | import time 4 | import urllib.parse 5 | from contextlib import contextmanager 6 | from dataclasses import dataclass 7 | from queue import Empty, Queue 8 | from threading import Thread 9 | from typing import Any, Dict, Generator, List, Optional, Tuple 10 | 11 | import requests 12 | from requests.adapters import HTTPAdapter, Retry 13 | 14 | from .check_type import CheckTypeException, OptionalKey, OptionalValue, check_type 15 | 16 | 17 | @dataclass 18 | class Task: 19 | """A task returned by the remote server.""" 20 | 21 | id: str 22 | contents: str 23 | 24 | 25 | @dataclass 26 | class QueueCounts: 27 | """Task counts for a single queue.""" 28 | 29 | pending: int 30 | running: int 31 | expired: int 32 | completed: int 33 | 34 | # Won't be set if a time window wasn't specified in the request, or if the 35 | # server is old enough to not support rate estimation. 36 | rate: Optional[float] = None 37 | 38 | modtime: Optional[int] = None 39 | bytes: Optional[int] = None 40 | 41 | 42 | class TasqClient: 43 | """ 44 | An object for interacting with a specific queue in a tasq server. 45 | 46 | :param base_url: the URL for the server, possibly with a trailing 47 | slash, but with no extra path appended to it. 48 | E.g. "http://tasq.mydomain.com/". 49 | :param keepalive_interval: for running tasks from pop_running_task(), 50 | this is how often keepalives are sent. 51 | :param context: the queue context (empty string uses default queue). 52 | :param username: optional username for basic authentication. 53 | :param password: optional password for basic authentication. 54 | :param max_timeout: the maximum amount of time (in seconds) to wait 55 | between attempts to pop a task in pop_running_task(), 56 | or push a task in push_blocking(). 57 | Lower values mean waiting less long to pop in the case 58 | that a new task is pushed or all tasks are finished. 59 | :param task_timeout: if specified, override the timeout on the server with 60 | a custom timeout. This can be useful if we know we 61 | will be sending frequent keepalives, but the server 62 | has a longer timeout period. 63 | :param retry_server_errors: if True, retry requests if the server returns 64 | certain 5xx status codes. 65 | """ 66 | 67 | def __init__( 68 | self, 69 | base_url: str, 70 | keepalive_interval: float = 30.0, 71 | context: str = "", 72 | username: Optional[str] = None, 73 | password: Optional[str] = None, 74 | max_timeout: float = 30.0, 75 | task_timeout: Optional[float] = None, 76 | retry_server_errors: bool = True, 77 | ): 78 | self.base_url = base_url.rstrip("/") 79 | self.keepalive_interval = keepalive_interval 80 | self.context = context 81 | self.username = username 82 | self.password = password 83 | self.max_timeout = max_timeout 84 | self.task_timeout = task_timeout 85 | self.retry_server_errors = retry_server_errors 86 | self.session = requests.Session() 87 | self._configure_session() 88 | 89 | def push(self, contents: str, limit: int = 0) -> Optional[str]: 90 | """ 91 | Push a task and get its resulting ID. 92 | 93 | If limit is specified, then the task will not be pushed if the queue is 94 | full, in which case None is returned. 95 | """ 96 | return self._post_form( 97 | "/task/push", dict(contents=contents, limit=limit), type_template=OptionalValue(str) 98 | ) 99 | 100 | def push_batch(self, ids: List[str], limit: int = 0) -> Optional[List[str]]: 101 | """ 102 | Push a batch of tasks and get their resulting IDs. 103 | 104 | If limit is specified, then tasks will not be pushed if the queue does 105 | not have room for all the tasks at once, in which case None is 106 | returned. 107 | 108 | If limit is negative, then (-limit + batch_size) is used as the limit. 109 | This effectively limits the size of the queue before a push rather than 110 | after the push, to prevent large batches from being less likely to be 111 | pushed than larger batches. 112 | """ 113 | if limit < 0: 114 | limit = -limit + len(ids) 115 | return self._post_json( 116 | f"/task/push_batch?limit={limit}", ids, type_template=OptionalValue([str]) 117 | ) 118 | 119 | def push_blocking( 120 | self, contents: List[str], limit: int, init_wait_time: float = 1.0 121 | ) -> List[str]: 122 | """ 123 | Push one or more tasks atomically and block until they are pushed. 124 | 125 | If the queue cannot fit the batch, this will wait to retry with random 126 | exponential backoff. Backoff is randomized to mitigate starvation. 127 | 128 | See push_batch() for details on passing a negative limit to avoid 129 | starvation of larger batches when pushing from multiple processes. 130 | 131 | Unlike push_batch(), the ids returned by this method will never be 132 | None, since all tasks must be pushed. 133 | """ 134 | assert isinstance(contents, (list, tuple)), ( 135 | f"expected a list of task contents, got object of type {type(contents)}" 136 | ) 137 | assert init_wait_time <= self.max_timeout, ( 138 | f"wait time {init_wait_time=} should not be larger than {self.max_timeout=}" 139 | ) 140 | assert limit < 0 or limit >= len(contents) 141 | 142 | cur_wait = init_wait_time 143 | while True: 144 | ids = self.push_batch(contents, limit=limit) 145 | if ids is not None: 146 | return ids 147 | timeout = cur_wait * random.random() 148 | time.sleep(timeout) 149 | # Use summation instead of doubling to prevent really rapid 150 | # growth of cur_wait with low probability. 151 | cur_wait = min(cur_wait + timeout, self.max_timeout) 152 | 153 | def pop(self) -> Tuple[Optional[Task], Optional[float]]: 154 | """ 155 | Pop a pending task from the queue. 156 | 157 | If no task is returned, a retry time may be returned indicating the 158 | number of seconds until the next in-progress task will expire. If this 159 | retry time is also None, then the queue has been exhausted. 160 | """ 161 | result = self._get( 162 | "/task/pop", 163 | type_template={ 164 | OptionalKey("id"): str, 165 | OptionalKey("contents"): str, 166 | OptionalKey("retry"): float, 167 | OptionalKey("done"): bool, 168 | }, 169 | supports_timeout=True, 170 | ) 171 | if "id" in result and "contents" in result: 172 | return Task(id=result["id"], contents=result["contents"]), None 173 | elif "done" not in result: 174 | raise TasqMisbehavingServerError("no done field in response") 175 | elif result["done"]: 176 | return None, None 177 | elif "retry" not in result: 178 | raise TasqMisbehavingServerError("missing retry value") 179 | else: 180 | return None, float(result["retry"]) 181 | 182 | def pop_batch(self, n: int) -> Tuple[List[Task], Optional[float]]: 183 | """ 184 | Retrieve at most n tasks from the queue. 185 | 186 | If fewer than n tasks are returned, a retry time may be returned to 187 | indicate when the next pending task will expire. 188 | 189 | If no tasks are returned and the retry time is None, then the queue has 190 | been exhausted. 191 | """ 192 | response = self._post_form( 193 | "/task/pop_batch", 194 | dict(count=n), 195 | type_template={ 196 | "done": bool, 197 | "tasks": [dict(id=str, contents=str)], 198 | OptionalKey("retry"): float, 199 | }, 200 | supports_timeout=True, 201 | ) 202 | 203 | if response["done"]: 204 | return [], None 205 | 206 | retry = float(response["retry"]) if "retry" in response else None 207 | 208 | if len(response["tasks"]): 209 | return [Task(id=x["id"], contents=x["contents"]) for x in response["tasks"]], retry 210 | elif retry is not None: 211 | return [], retry 212 | else: 213 | raise TasqMisbehavingServerError( 214 | "no retry time specified when tasks are empty and done is false" 215 | ) 216 | 217 | def completed(self, id: str): 218 | """Indicate that an in-progress task has been completed.""" 219 | self._post_form("/task/completed", dict(id=id)) 220 | 221 | def completed_batch(self, ids: List[str]): 222 | """Indicate that some in-progress tasks have been completed.""" 223 | self._post_json("/task/completed_batch", ids) 224 | 225 | def keepalive(self, id: str): 226 | """Reset the timeout interval for a still in-progress task.""" 227 | self._post_form("/task/keepalive", dict(id=id), supports_timeout=True) 228 | 229 | @contextmanager 230 | def pop_running_task(self) -> Generator[Optional["RunningTask"], None, None]: 231 | """ 232 | Pop a task from the queue and wrap it in a RunningTask, blocking until 233 | the queue is completely empty or a task is successfully popped. 234 | 235 | The resulting RunningTask manages a background process that sends 236 | keepalives for the returned task ID at regular intervals. 237 | 238 | This is meant to be used in a `with` clause. When the `with` clause is 239 | exited, the keepalive loop is stopped, and the task will be marked as 240 | completed unless the with clause is exited with an exception. 241 | """ 242 | while True: 243 | task, timeout = self.pop() 244 | if task is not None: 245 | rt = RunningTask(self, id=task.id, contents=task.contents) 246 | try: 247 | yield rt 248 | rt.completed() 249 | finally: 250 | rt.cancel() 251 | return 252 | elif timeout is not None: 253 | time.sleep(min(timeout, self.max_timeout)) 254 | else: 255 | yield None 256 | return 257 | 258 | def counts(self, rate_window: int = 0) -> QueueCounts: 259 | """Get the number of tasks in each state within the queue.""" 260 | data = self._get( 261 | f"/counts?window={rate_window}&includeModtime=1&includeBytes=1", 262 | { 263 | "pending": int, 264 | "running": int, 265 | "expired": int, 266 | "completed": int, 267 | OptionalKey("minute_rate"): float, 268 | OptionalKey("modtime"): int, 269 | OptionalKey("bytes"): int, 270 | }, 271 | ) 272 | return QueueCounts(**data) 273 | 274 | def clear(self): 275 | """Deletes the queue and all tasks in it.""" 276 | result = self._post_form("/task/clear", dict()) 277 | if result is True: 278 | return True 279 | else: 280 | raise TasqMisbehavingServerError("failed to clear queue") 281 | 282 | def __getstate__( 283 | self, 284 | ): 285 | res = self.__dict__.copy() 286 | del res["session"] 287 | return res 288 | 289 | def __setstate__(self, state: Dict[str, Any]): 290 | self.__dict__ = state 291 | self.session = requests.Session() 292 | self._configure_session() 293 | 294 | def _configure_session(self): 295 | if self.username is not None or self.password is not None: 296 | assert self.username is not None and self.password is not None 297 | self.session.auth = (self.username, self.password) 298 | if self.retry_server_errors: 299 | retries = Retry( 300 | total=10, 301 | backoff_factor=1.0, 302 | status_forcelist=[500, 502, 503, 504], 303 | allowed_methods=False, 304 | ) 305 | for schema in ("http://", "https://"): 306 | self.session.mount(schema, HTTPAdapter(max_retries=retries)) 307 | 308 | def _get( 309 | self, path: str, type_template: Optional[Any] = None, supports_timeout: bool = False 310 | ) -> Any: 311 | return _process_response( 312 | self.session.get(self._url_for_path(path, supports_timeout)), type_template 313 | ) 314 | 315 | def _post_form( 316 | self, 317 | path: str, 318 | args: Dict[str, str], 319 | type_template: Optional[Any] = None, 320 | supports_timeout: bool = False, 321 | ) -> Any: 322 | return _process_response( 323 | self.session.post(self._url_for_path(path, supports_timeout), data=args), type_template 324 | ) 325 | 326 | def _post_json( 327 | self, 328 | path: str, 329 | data: Any, 330 | type_template: Optional[Any] = None, 331 | supports_timeout: bool = False, 332 | ) -> Any: 333 | return _process_response( 334 | self.session.post(self._url_for_path(path, supports_timeout), json=data), type_template 335 | ) 336 | 337 | def _url_for_path(self, path: str, supports_timeout: bool) -> str: 338 | separator = "?" if "?" not in path else "&" 339 | result = self.base_url + path + separator + "context=" + urllib.parse.quote(self.context) 340 | if supports_timeout and self.task_timeout is not None: 341 | result += "&timeout=" + urllib.parse.quote(f"{self.task_timeout:f}") 342 | return result 343 | 344 | 345 | @dataclass 346 | class RunningTask(Task): 347 | """ 348 | A task object that periodically sends keepalives in the background until 349 | cancel() or completed() is called. 350 | """ 351 | 352 | def __init__(self, client: TasqClient, *args, **kwargs): 353 | super().__init__(*args, **kwargs) 354 | self.client = client 355 | self._kill_queue = Queue() 356 | self._thread = Thread( 357 | target=RunningTask._keepalive_worker, 358 | name="tasq-keepalive-worker", 359 | args=( 360 | self._kill_queue, 361 | client, 362 | self.id, 363 | ), 364 | daemon=True, 365 | ) 366 | self._thread.start() 367 | 368 | def cancel(self): 369 | if self._thread is None: 370 | return 371 | self._kill_queue.put(None) 372 | self._thread.join() 373 | self._thread = None 374 | 375 | def completed(self): 376 | self.cancel() 377 | self.client.completed(self.id) 378 | 379 | @staticmethod 380 | def _keepalive_worker( 381 | kill_queue: Queue, 382 | client: TasqClient, 383 | task_id: str, 384 | ): 385 | while True: 386 | try: 387 | client.keepalive(task_id) 388 | except Exception as exc: # pylint: disable=broad-except 389 | # Ignore the error if we killed the thread during the 390 | # keepalive call. 391 | try: 392 | kill_queue.get(block=False) 393 | return 394 | except Empty: 395 | pass 396 | print(f"exception in tasq keepalive worker: {exc}", file=sys.stderr) 397 | try: 398 | kill_queue.get(timeout=client.keepalive_interval) 399 | return 400 | except Empty: 401 | pass 402 | 403 | 404 | class TasqRemoteError(Exception): 405 | """An error returned by a remote server.""" 406 | 407 | 408 | class TasqMisbehavingServerError(Exception): 409 | """An error when a tasq server misbehaves.""" 410 | 411 | 412 | def _process_response(response: requests.Response, type_template: Optional[Any]) -> Any: 413 | try: 414 | parsed = response.json() 415 | except Exception as exc: 416 | raise TasqMisbehavingServerError("failed to get JSON from response") from exc 417 | 418 | check_template = { 419 | OptionalKey("error"): str, 420 | OptionalKey("data"): object if type_template is None else type_template, 421 | } 422 | try: 423 | check_type(check_template, parsed) 424 | except CheckTypeException as exc: 425 | raise TasqMisbehavingServerError(f"invalid response object: {exc}") from exc 426 | 427 | if "error" in parsed: 428 | raise TasqRemoteError(parsed["error"]) 429 | elif "data" in parsed: 430 | return parsed["data"] 431 | else: 432 | raise TasqMisbehavingServerError("missing error or data fields in response") 433 | -------------------------------------------------------------------------------- /tasq-rate-estimate/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "log" 6 | "time" 7 | 8 | "github.com/unixpickle/essentials" 9 | "github.com/unixpickle/tasq" 10 | ) 11 | 12 | func main() { 13 | var host string 14 | var context string 15 | var username string 16 | var password string 17 | var interval time.Duration 18 | flag.StringVar(&host, "host", "", "server URL") 19 | flag.StringVar(&context, "context", "", "tasq context name") 20 | flag.StringVar(&username, "username", "", "basic auth username") 21 | flag.StringVar(&password, "password", "", "basic auth password") 22 | flag.DurationVar(&interval, "interval", time.Second, "number of seconds between count calls") 23 | flag.Parse() 24 | 25 | if host == "" { 26 | essentials.Die("Must provide -host argument. See -help.") 27 | } 28 | 29 | client, err := tasq.NewClient(host, context, username, password) 30 | essentials.Must(err) 31 | 32 | t1 := time.Now() 33 | startCounts, err := client.QueueCounts() 34 | essentials.Must(err) 35 | 36 | for { 37 | time.Sleep(interval) 38 | counts, err := client.QueueCounts() 39 | essentials.Must(err) 40 | completed := float64(counts.Completed - startCounts.Completed) 41 | elapsed := time.Now().Sub(t1).Seconds() 42 | log.Printf("task rate: %.03f tasks/second (total time %.02f seconds)", completed/elapsed, elapsed) 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /tasq-server/homepage.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | const Homepage = ` 4 | 5 | 6 | 7 | 257 | 258 | 259 |
260 |
261 | 262 |
263 |
264 | 270 |
271 |
272 | 273 | 274 |
275 |
276 |
    277 | 280 | 281 |
    282 |

    Quickly add a task

    283 |
    284 | 285 | 286 |
    287 |
    288 | 289 | 290 |
    291 | 292 |
    293 |
  1. 294 | 295 | 296 | 297 | 298 | 299 | 300 | 301 | 302 | 303 | 304 | 305 | 306 | 307 | 308 | 309 | 310 | 311 | 312 | 313 | 314 | 315 | 316 | 317 | 318 | 319 | 320 | 321 | 322 | 323 | 324 |
    Uptime:-
    Allocated:-
    Total allocated:-
    System allocated:-
    Last GC:-
    Last save:-
    Save latency:-
    325 |
    2. 326 |
    327 |
    328 | 329 | 330 |
    331 |
    332 | 333 | 711 | 712 | 713 | ` 714 | -------------------------------------------------------------------------------- /tasq-server/json.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "io" 7 | ) 8 | 9 | type JSONWriter interface { 10 | WriteJSON(w io.Writer) error 11 | } 12 | 13 | func WriteJSONObject(w io.Writer, obj map[string]interface{}) error { 14 | if _, err := w.Write([]byte("{")); err != nil { 15 | return err 16 | } 17 | first := true 18 | for k, v := range obj { 19 | if first { 20 | first = false 21 | } else { 22 | if _, err := w.Write([]byte(",")); err != nil { 23 | return err 24 | } 25 | } 26 | 27 | // There are probably some extra restrictions on JSON keys, but 28 | // we ignore these for now. 29 | encKey, err := json.Marshal(k) 30 | if err != nil { 31 | return err 32 | } 33 | if _, err := w.Write(append(encKey, ':')); err != nil { 34 | return err 35 | } 36 | 37 | if writer, ok := v.(JSONWriter); ok { 38 | if err := writer.WriteJSON(w); err != nil { 39 | return err 40 | } 41 | } else { 42 | encoded, err := json.Marshal(v) 43 | if err != nil { 44 | return err 45 | } 46 | if _, err := w.Write(encoded); err != nil { 47 | return err 48 | } 49 | } 50 | } 51 | if _, err := w.Write([]byte("}")); err != nil { 52 | return err 53 | } 54 | return nil 55 | } 56 | 57 | type EncodedTaskList []EncodedTask 58 | 59 | func (e EncodedTaskList) WriteJSON(w io.Writer) error { 60 | encodedStream := make(chan []byte, 32) 61 | ctx, cancel := context.WithCancel(context.Background()) 62 | defer cancel() 63 | 64 | go func() { 65 | defer close(encodedStream) 66 | for _, t := range e { 67 | data, err := json.Marshal(t) 68 | if err != nil { 69 | panic(err) 70 | } 71 | select { 72 | case encodedStream <- data: 73 | case <-ctx.Done(): 74 | return 75 | } 76 | } 77 | }() 78 | 79 | first := true 80 | if _, err := w.Write([]byte("[")); err != nil { 81 | return err 82 | } 83 | for encoded := range encodedStream { 84 | if first { 85 | first = false 86 | } else { 87 | if _, err := w.Write([]byte(",")); err != nil { 88 | return err 89 | } 90 | } 91 | if _, err := w.Write(encoded); err != nil { 92 | return err 93 | } 94 | } 95 | if _, err := w.Write([]byte("]")); err != nil { 96 | return err 97 | } 98 | return nil 99 | } 100 | -------------------------------------------------------------------------------- /tasq-server/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "crypto/subtle" 6 | "encoding/json" 7 | "errors" 8 | "flag" 9 | "fmt" 10 | "io" 11 | "log" 12 | "math" 13 | "net/http" 14 | "os" 15 | "os/signal" 16 | "runtime" 17 | "strconv" 18 | "strings" 19 | "sync" 20 | "syscall" 21 | "time" 22 | 23 | "github.com/unixpickle/essentials" 24 | ) 25 | 26 | func main() { 27 | var addr string 28 | var pathPrefix string 29 | var authUsername string 30 | var authPassword string 31 | var savePath string 32 | var saveInterval time.Duration 33 | var timeout time.Duration 34 | flag.StringVar(&addr, "addr", ":8080", "address to listen on") 35 | flag.StringVar(&pathPrefix, "path-prefix", "/", "prefix for URL paths") 36 | flag.StringVar(&authUsername, "auth-username", "", "username for basic auth") 37 | flag.StringVar(&authPassword, "auth-password", "", "password for basic auth") 38 | flag.StringVar(&savePath, "save-path", "", "if specified, path to periodically save state to") 39 | flag.DurationVar(&timeout, "timeout", time.Minute*15, "timeout of individual tasks") 40 | flag.DurationVar(&saveInterval, "save-interval", time.Minute*5, "time between saves") 41 | flag.Parse() 42 | 43 | if !strings.HasSuffix(pathPrefix, "/") || !strings.HasPrefix(pathPrefix, "/") { 44 | essentials.Die("path prefix must start and end with a '/' character") 45 | } 46 | 47 | s := &Server{ 48 | PathPrefix: pathPrefix, 49 | AuthUsername: authUsername, 50 | AuthPassword: authPassword, 51 | SavePath: savePath, 52 | SaveInterval: saveInterval, 53 | StartTime: time.Now(), 54 | Queues: NewQueueStateMux(timeout), 55 | } 56 | http.HandleFunc(pathPrefix, s.ServeIndex) 57 | http.HandleFunc(pathPrefix+"summary", s.ServeSummary) 58 | http.HandleFunc(pathPrefix+"counts", s.ServeCounts) 59 | http.HandleFunc(pathPrefix+"stats", s.ServeStats) 60 | http.HandleFunc(pathPrefix+"task/push", s.ServePushTask) 61 | http.HandleFunc(pathPrefix+"task/push_batch", s.ServePushBatch) 62 | http.HandleFunc(pathPrefix+"task/pop", s.ServePopTask) 63 | http.HandleFunc(pathPrefix+"task/pop_batch", s.ServePopBatch) 64 | http.HandleFunc(pathPrefix+"task/peek", s.ServePeekTask) 65 | http.HandleFunc(pathPrefix+"task/completed", s.ServeCompletedTask) 66 | http.HandleFunc(pathPrefix+"task/completed_batch", s.ServeCompletedBatch) 67 | http.HandleFunc(pathPrefix+"task/keepalive", s.ServeKeepalive) 68 | http.HandleFunc(pathPrefix+"task/clear", s.ServeClearTasks) 69 | http.HandleFunc(pathPrefix+"task/expire_all", s.ServeExpireTasks) 70 | http.HandleFunc(pathPrefix+"task/queue_expired", s.ServeQueueExpired) 71 | s.SetupSaveLoop(timeout) 72 | essentials.Must(http.ListenAndServe(addr, nil)) 73 | } 74 | 75 | type Server struct { 76 | PathPrefix string 77 | AuthUsername string 78 | AuthPassword string 79 | Queues *QueueStateMux 80 | SavePath string 81 | SaveInterval time.Duration 82 | 83 | StartTime time.Time 84 | 85 | SaveStatsLock sync.RWMutex 86 | LastSave time.Time 87 | LastSaveDuration time.Duration 88 | 89 | SignalChan <-chan os.Signal 90 | } 91 | 92 | func (s *Server) ServeIndex(w http.ResponseWriter, r *http.Request) { 93 | if !s.BasicAuth(w, r) { 94 | return 95 | } 96 | if r.URL.Path == s.PathPrefix || r.URL.Path+"/" == s.PathPrefix { 97 | w.Header().Set("content-type", "text/html") 98 | w.Write([]byte(Homepage)) 99 | } else { 100 | w.Header().Set("content-type", "text/html") 101 | w.WriteHeader(http.StatusNotFound) 102 | fmt.Fprintln(w, "Page not found") 103 | } 104 | } 105 | 106 | func (s *Server) ServeSummary(w http.ResponseWriter, r *http.Request) { 107 | if !s.BasicAuth(w, r) { 108 | return 109 | } 110 | w.Header().Set("content-type", "text/plain") 111 | found := false 112 | buf := bytes.NewBuffer(nil) 113 | err := s.Queues.Iterate(func(name string, qs *QueueState) { 114 | found = true 115 | if name == "" { 116 | fmt.Fprint(buf, "---- Default context ----\n") 117 | } else { 118 | fmt.Fprintf(buf, "---- Context: %s ----\n", name) 119 | } 120 | counts := qs.Counts(0, false, true) 121 | fmt.Fprintf(buf, " Pending: %d\n", counts.Pending) 122 | fmt.Fprintf(buf, "In progress: %d\n", counts.Running) 123 | fmt.Fprintf(buf, " Expired: %d\n", counts.Expired) 124 | fmt.Fprintf(buf, " Completed: %d\n", counts.Completed) 125 | fmt.Fprintf(buf, " Bytes: %d\n", counts.Bytes) 126 | }) 127 | if err != nil { 128 | fmt.Fprint(buf, err.Error()) 129 | w.WriteHeader(http.StatusServiceUnavailable) 130 | } else if !found { 131 | fmt.Fprint(buf, "No active queues.") 132 | } 133 | w.Write(buf.Bytes()) 134 | } 135 | 136 | func (s *Server) ServeCounts(w http.ResponseWriter, r *http.Request) { 137 | if !s.BasicAuth(w, r) { 138 | return 139 | } 140 | 141 | var rateWindow int 142 | if s := r.URL.Query().Get("window"); s != "" { 143 | var err error 144 | rateWindow, err = strconv.Atoi(s) 145 | if err != nil { 146 | serveError(w, err.Error(), http.StatusBadRequest) 147 | return 148 | } 149 | } 150 | 151 | includeModtime := r.URL.Query().Get("includeModtime") == "1" 152 | includeBytes := r.URL.Query().Get("includeBytes") == "1" 153 | 154 | if r.URL.Query().Get("all") == "1" { 155 | allNames := []string{} 156 | allCounts := []*QueueCounts{} 157 | err := s.Queues.Iterate(func(name string, qs *QueueState) { 158 | allNames = append(allNames, name) 159 | allCounts = append(allCounts, qs.Counts(rateWindow, includeModtime, includeBytes)) 160 | }) 161 | if err != nil { 162 | serveError(w, err.Error(), http.StatusServiceUnavailable) 163 | } else { 164 | serveObject(w, map[string]interface{}{ 165 | "names": allNames, 166 | "counts": allCounts, 167 | }) 168 | } 169 | return 170 | } 171 | 172 | var obj interface{} 173 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 174 | obj = qs.Counts(rateWindow, includeModtime, includeBytes) 175 | }) 176 | if err != nil { 177 | serveError(w, err.Error(), http.StatusServiceUnavailable) 178 | } else { 179 | serveObject(w, obj) 180 | } 181 | } 182 | 183 | func (s *Server) ServeStats(w http.ResponseWriter, r *http.Request) { 184 | if !s.BasicAuth(w, r) { 185 | return 186 | } 187 | 188 | var m runtime.MemStats 189 | runtime.ReadMemStats(&m) 190 | 191 | s.SaveStatsLock.RLock() 192 | saveStats := map[string]interface{}{ 193 | "elapsed": (time.Since(s.LastSave).Seconds()), 194 | "latency": s.LastSaveDuration.Seconds(), 195 | } 196 | s.SaveStatsLock.RUnlock() 197 | 198 | serveObject(w, map[string]interface{}{ 199 | "uptime": time.Since(s.StartTime).Seconds(), 200 | "memory": map[string]interface{}{ 201 | "alloc": m.Alloc, 202 | "totalAlloc": m.TotalAlloc, 203 | "sys": m.Sys, 204 | "lastGC": float64(time.Now().UnixNano()-int64(m.LastGC)) / 1000000000.0, 205 | }, 206 | "save": saveStats, 207 | }) 208 | } 209 | 210 | func (s *Server) ServePushTask(w http.ResponseWriter, r *http.Request) { 211 | if !s.BasicAuth(w, r) { 212 | return 213 | } 214 | contents := r.FormValue("contents") 215 | limit, err := parseLimit(r.FormValue("limit")) 216 | if err != nil { 217 | serveError(w, err.Error(), http.StatusBadRequest) 218 | return 219 | } 220 | if contents == "" { 221 | serveError(w, "must specify non-empty `contents` parameter", http.StatusBadRequest) 222 | } else { 223 | var obj interface{} 224 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 225 | if id, ok := qs.Push(contents, limit); ok { 226 | obj = id 227 | } 228 | }) 229 | if err != nil { 230 | serveError(w, err.Error(), http.StatusServiceUnavailable) 231 | } else { 232 | serveObject(w, obj) 233 | } 234 | } 235 | } 236 | 237 | func (s *Server) ServePushBatch(w http.ResponseWriter, r *http.Request) { 238 | if !s.BasicAuth(w, r) { 239 | return 240 | } 241 | data, err := io.ReadAll(r.Body) 242 | if err != nil { 243 | return 244 | } 245 | var contents []string 246 | if err := json.Unmarshal(data, &contents); err != nil { 247 | serveError(w, err.Error(), http.StatusBadRequest) 248 | } else { 249 | limit, err := parseLimit(r.URL.Query().Get("limit")) 250 | if err != nil { 251 | serveError(w, err.Error(), http.StatusBadRequest) 252 | return 253 | } 254 | var ids []string 255 | err = s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 256 | ids, _ = qs.PushBatch(contents, limit) 257 | }) 258 | if err != nil { 259 | serveError(w, err.Error(), http.StatusServiceUnavailable) 260 | } else { 261 | serveObject(w, ids) 262 | } 263 | } 264 | } 265 | 266 | func (s *Server) ServePopTask(w http.ResponseWriter, r *http.Request) { 267 | if !s.BasicAuth(w, r) { 268 | return 269 | } 270 | timeout, timeoutOk := s.TimeoutParam(w, r) 271 | if !timeoutOk { 272 | return 273 | } 274 | 275 | var task *Task 276 | var nextTry *time.Time 277 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 278 | task, nextTry = qs.Pop(timeout) 279 | }) 280 | if err != nil { 281 | serveError(w, err.Error(), http.StatusServiceUnavailable) 282 | } else if task != nil { 283 | serveObject(w, task) 284 | } else { 285 | if nextTry != nil { 286 | timeout := time.Until(*nextTry) 287 | serveObject(w, map[string]interface{}{ 288 | "done": false, 289 | "retry": math.Max(0, timeout.Seconds()), 290 | }) 291 | } else { 292 | serveObject(w, map[string]interface{}{"done": true}) 293 | } 294 | } 295 | } 296 | 297 | func (s *Server) ServePopBatch(w http.ResponseWriter, r *http.Request) { 298 | if !s.BasicAuth(w, r) { 299 | return 300 | } 301 | timeout, timeoutOk := s.TimeoutParam(w, r) 302 | if !timeoutOk { 303 | return 304 | } 305 | 306 | n, err := strconv.Atoi(r.FormValue("count")) 307 | if err != nil { 308 | serveError(w, "invalid 'count' parameter: "+err.Error(), http.StatusBadRequest) 309 | return 310 | } else if n <= 0 { 311 | serveError(w, "invalid 'count' requested", http.StatusBadRequest) 312 | return 313 | } 314 | 315 | var tasks []*Task 316 | var nextTry *time.Time 317 | err = s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 318 | tasks, nextTry = qs.PopBatch(n, timeout) 319 | }) 320 | if err != nil { 321 | serveError(w, err.Error(), http.StatusServiceUnavailable) 322 | return 323 | } 324 | 325 | result := map[string]interface{}{ 326 | "done": len(tasks) == 0 && nextTry == nil, 327 | } 328 | if nextTry != nil { 329 | timeout := time.Until(*nextTry) 330 | result["retry"] = math.Max(0, timeout.Seconds()) 331 | } 332 | if tasks == nil { 333 | // Prevent a null value in the JSON field. 334 | tasks = []*Task{} 335 | } 336 | result["tasks"] = tasks 337 | 338 | serveObject(w, result) 339 | } 340 | 341 | func (s *Server) ServePeekTask(w http.ResponseWriter, r *http.Request) { 342 | if !s.BasicAuth(w, r) { 343 | return 344 | } 345 | var task, nextTask *Task 346 | var nextTime *time.Time 347 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 348 | task, nextTask, nextTime = qs.Peek() 349 | }) 350 | if err != nil { 351 | serveError(w, err.Error(), http.StatusServiceUnavailable) 352 | } else if task != nil { 353 | serveObject(w, map[string]interface{}{"contents": task.Contents, "id": task.ID}) 354 | } else { 355 | if nextTask != nil { 356 | timeout := time.Until(*nextTime) 357 | serveObject(w, map[string]interface{}{ 358 | "done": false, 359 | "retry": math.Max(0, timeout.Seconds()), 360 | "next": map[string]interface{}{ 361 | "contents": nextTask.Contents, 362 | "id": nextTask.ID, 363 | }, 364 | }) 365 | } else { 366 | serveObject(w, map[string]interface{}{"done": true}) 367 | } 368 | } 369 | } 370 | 371 | func (s *Server) ServeCompletedTask(w http.ResponseWriter, r *http.Request) { 372 | if !s.BasicAuth(w, r) { 373 | return 374 | } 375 | id := r.FormValue("id") 376 | var status bool 377 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 378 | status = qs.Completed(id) 379 | }) 380 | if err != nil { 381 | serveError(w, err.Error(), http.StatusServiceUnavailable) 382 | } else if status { 383 | serveObject(w, true) 384 | } else { 385 | serveError(w, "there was no in-progress task with the specified `id`", http.StatusOK) 386 | } 387 | } 388 | 389 | func (s *Server) ServeCompletedBatch(w http.ResponseWriter, r *http.Request) { 390 | if !s.BasicAuth(w, r) { 391 | return 392 | } 393 | data, err := io.ReadAll(r.Body) 394 | if err != nil { 395 | return 396 | } 397 | var ids []string 398 | if err := json.Unmarshal(data, &ids); err != nil { 399 | serveError(w, err.Error(), http.StatusBadRequest) 400 | } else { 401 | var failures []string 402 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 403 | for _, id := range ids { 404 | if !qs.Completed(id) { 405 | failures = append(failures, id) 406 | } 407 | } 408 | }) 409 | if err != nil { 410 | serveError(w, err.Error(), http.StatusServiceUnavailable) 411 | return 412 | } 413 | if len(failures) > 0 { 414 | serveError(w, "there were no in-progress tasks with the specified ids: "+ 415 | strings.Join(failures, ", "), http.StatusOK) 416 | } else { 417 | serveObject(w, true) 418 | } 419 | } 420 | } 421 | 422 | func (s *Server) ServeKeepalive(w http.ResponseWriter, r *http.Request) { 423 | if !s.BasicAuth(w, r) { 424 | return 425 | } 426 | timeout, timeoutOk := s.TimeoutParam(w, r) 427 | if !timeoutOk { 428 | return 429 | } 430 | id := r.FormValue("id") 431 | 432 | var status bool 433 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 434 | status = qs.Keepalive(id, timeout) 435 | }) 436 | if err != nil { 437 | serveError(w, err.Error(), http.StatusServiceUnavailable) 438 | } else if status { 439 | serveObject(w, true) 440 | } else { 441 | serveError(w, "there was no in-progress task with the specified `id`", http.StatusOK) 442 | } 443 | } 444 | 445 | func (s *Server) ServeClearTasks(w http.ResponseWriter, r *http.Request) { 446 | if !s.BasicAuth(w, r) { 447 | return 448 | } 449 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 450 | qs.Clear() 451 | }) 452 | if err != nil { 453 | serveError(w, err.Error(), http.StatusServiceUnavailable) 454 | } else { 455 | serveObject(w, true) 456 | } 457 | } 458 | 459 | func (s *Server) ServeExpireTasks(w http.ResponseWriter, r *http.Request) { 460 | if !s.BasicAuth(w, r) { 461 | return 462 | } 463 | var n int 464 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 465 | n = qs.ExpireAll() 466 | }) 467 | if err != nil { 468 | serveError(w, err.Error(), http.StatusServiceUnavailable) 469 | } else { 470 | serveObject(w, n) 471 | } 472 | } 473 | 474 | func (s *Server) ServeQueueExpired(w http.ResponseWriter, r *http.Request) { 475 | if !s.BasicAuth(w, r) { 476 | return 477 | } 478 | var n int 479 | err := s.Queues.Get(r.URL.Query().Get("context"), func(qs *QueueState) { 480 | n = qs.QueueExpired() 481 | }) 482 | if err != nil { 483 | serveError(w, err.Error(), http.StatusServiceUnavailable) 484 | } else { 485 | serveObject(w, n) 486 | } 487 | } 488 | 489 | func (s *Server) BasicAuth(w http.ResponseWriter, r *http.Request) bool { 490 | if s.AuthUsername == "" && s.AuthPassword == "" { 491 | return true 492 | } 493 | username, password, ok := r.BasicAuth() 494 | if !ok { 495 | w.Header().Set("www-authenticate", `Basic realm="restricted", charset="UTF-8"`) 496 | w.Header().Set("content-type", "application/json") 497 | w.WriteHeader(http.StatusUnauthorized) 498 | w.Write([]byte(`{"error": "basic auth must be provided"}`)) 499 | return false 500 | } 501 | if subtle.ConstantTimeCompare([]byte(username), []byte(s.AuthUsername)) == 1 && 502 | subtle.ConstantTimeCompare([]byte(password), []byte(s.AuthPassword)) == 1 { 503 | return true 504 | } else { 505 | w.Header().Set("www-authenticate", `Basic realm="restricted", charset="UTF-8"`) 506 | w.Header().Set("content-type", "application/json") 507 | w.WriteHeader(http.StatusUnauthorized) 508 | w.Write([]byte(`{"error": "incorrect credentials"}`)) 509 | return false 510 | } 511 | } 512 | 513 | func (s *Server) TimeoutParam(w http.ResponseWriter, r *http.Request) (*time.Duration, bool) { 514 | timeoutStr := r.URL.Query().Get("timeout") 515 | if timeoutStr == "" { 516 | return nil, true 517 | } 518 | parsed, err := strconv.ParseFloat(timeoutStr, 64) 519 | duration := time.Millisecond * time.Duration(parsed*1000) 520 | if err == nil && duration <= 0.0 { 521 | err = errors.New("timeout must be at least one millisecond") 522 | } 523 | if err != nil { 524 | w.Header().Set("www-authenticate", `Basic realm="restricted", charset="UTF-8"`) 525 | w.Header().Set("content-type", "application/json") 526 | w.WriteHeader(http.StatusUnauthorized) 527 | data, _ := json.Marshal(map[string]string{"error": err.Error()}) 528 | w.Write(data) 529 | return nil, false 530 | } 531 | return &duration, true 532 | } 533 | 534 | func (s *Server) SetupSaveLoop(timeout time.Duration) { 535 | if s.SavePath == "" { 536 | return 537 | } 538 | 539 | sigChan := make(chan os.Signal, 1) 540 | signal.Notify(sigChan, syscall.SIGUSR1) 541 | s.SignalChan = sigChan 542 | 543 | if _, err := os.Stat(s.SavePath); err == nil { 544 | log.Printf("Loading state from: %s", s.SavePath) 545 | s.Queues, err = ReadQueueStateMux(timeout, s.SavePath) 546 | if err != nil { 547 | log.Fatal(err) 548 | } else { 549 | log.Printf("Loaded state from: %s", s.SavePath) 550 | } 551 | } 552 | s.LastSave = time.Now() 553 | s.LastSaveDuration = 0 554 | go s.SaveLoop() 555 | } 556 | 557 | func (s *Server) SaveLoop() { 558 | var shutdown bool 559 | for !shutdown { 560 | select { 561 | case <-time.After(s.SaveInterval): 562 | case <-s.SignalChan: 563 | log.Println("caught SIGUSR1") 564 | shutdown = true 565 | } 566 | log.Printf("Saving state to: %s", s.SavePath) 567 | tmpPath := s.SavePath + ".tmp" 568 | w, err := os.Create(tmpPath) 569 | if err != nil { 570 | log.Fatal(err) 571 | } 572 | t1 := time.Now() 573 | err = s.Queues.Serialize(w, shutdown) 574 | w.Close() 575 | if err != nil { 576 | log.Fatal(err) 577 | } 578 | os.Rename(tmpPath, s.SavePath) 579 | 580 | s.SaveStatsLock.Lock() 581 | s.LastSave = time.Now() 582 | s.LastSaveDuration = s.LastSave.Sub(t1) 583 | s.SaveStatsLock.Unlock() 584 | 585 | log.Printf("Saved state to: %s", s.SavePath) 586 | } 587 | 588 | // We are shutting down post-save 589 | log.Println("exiting due to shutdown signal") 590 | os.Exit(0) 591 | } 592 | 593 | func parseLimit(limit string) (int, error) { 594 | if limit == "" { 595 | return 0, nil 596 | } 597 | value, err := strconv.Atoi(limit) 598 | if err != nil { 599 | return 0, err 600 | } 601 | return value, nil 602 | } 603 | 604 | func serveObject(w http.ResponseWriter, obj interface{}) { 605 | w.Header().Set("content-type", "application/json") 606 | json.NewEncoder(w).Encode(map[string]interface{}{"data": obj}) 607 | } 608 | 609 | func serveError(w http.ResponseWriter, err string, status int) { 610 | w.Header().Set("content-type", "application/json") 611 | w.WriteHeader(status) 612 | json.NewEncoder(w).Encode(map[string]interface{}{"error": err}) 613 | } 614 | -------------------------------------------------------------------------------- /tasq-server/queues.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "archive/zip" 5 | "bufio" 6 | "encoding/json" 7 | "io" 8 | "os" 9 | "sort" 10 | "strconv" 11 | "sync" 12 | "time" 13 | 14 | "github.com/pkg/errors" 15 | "github.com/unixpickle/essentials" 16 | ) 17 | 18 | var ErrShuttingDown = errors.New("server is shutting down") 19 | 20 | // QueueStateMux manages multiple (named) QueueStates. 21 | type QueueStateMux struct { 22 | saveLock sync.RWMutex 23 | finalError error 24 | lock sync.Mutex 25 | queues map[string]*QueueState 26 | users map[string]int 27 | timeout time.Duration 28 | } 29 | 30 | // NewQueueStateMux creates a QueueStateMux with the given task timeout. 31 | func NewQueueStateMux(timeout time.Duration) *QueueStateMux { 32 | return &QueueStateMux{ 33 | queues: map[string]*QueueState{}, 34 | users: map[string]int{}, 35 | timeout: timeout, 36 | } 37 | } 38 | 39 | // DeserializeQueueStateMux reads a file written by QueueStateMux.Serialize(). 40 | func DeserializeQueueStateMux(timeout time.Duration, r io.ReaderAt, 41 | size int64) (*QueueStateMux, error) { 42 | const context = "deserialize queue state" 43 | res := NewQueueStateMux(timeout) 44 | 45 | zf, err := zip.NewReader(r, size) 46 | if err != nil { 47 | return nil, errors.Wrap(err, context) 48 | } 49 | for _, file := range zf.File { 50 | subReader, err := file.Open() 51 | if err != nil { 52 | return nil, errors.Wrap(err, context) 53 | } 54 | var dictObj ContextState 55 | err = json.NewDecoder(subReader).Decode(&dictObj) 56 | subReader.Close() 57 | if err != nil { 58 | subReader.Close() 59 | return nil, errors.Wrap(err, context) 60 | } 61 | res.queues[dictObj.Name] = DecodeQueueState(dictObj.Encoded) 62 | res.users[dictObj.Name] = 0 63 | } 64 | return res, nil 65 | } 66 | 67 | // ReadQueueStateMux is like DeserializeQueueStateMux(), but reads from a local 68 | // file instead of an arbitrary reader. 69 | func ReadQueueStateMux(timeout time.Duration, path string) (*QueueStateMux, error) { 70 | stat, err := os.Stat(path) 71 | if err != nil { 72 | return nil, err 73 | } 74 | 75 | r, err := os.Open(path) 76 | if err != nil { 77 | return nil, err 78 | } 79 | defer r.Close() 80 | 81 | return DeserializeQueueStateMux(timeout, r, stat.Size()) 82 | } 83 | 84 | // Get calls f with a QueueState for the given name. One is created if 85 | // necessary, and will be destroyed when the queue is cleared. 86 | // 87 | // The QueueState should not be accessed outside of f. In particular, f should 88 | // not store a reference to the QueueState anywhere outside of its scope. 89 | // 90 | // Returns an error if the queue state has finished a final save. 91 | func (q *QueueStateMux) Get(name string, f func(*QueueState)) error { 92 | q.saveLock.RLock() 93 | defer q.saveLock.RUnlock() 94 | if q.finalError != nil { 95 | return q.finalError 96 | } 97 | q.get(name, f) 98 | return nil 99 | } 100 | 101 | func (q *QueueStateMux) get(name string, f func(*QueueState)) { 102 | q.lock.Lock() 103 | qs, ok := q.queues[name] 104 | if !ok { 105 | qs = NewQueueState(q.timeout) 106 | q.queues[name] = qs 107 | } 108 | q.users[name]++ 109 | q.lock.Unlock() 110 | 111 | defer func() { 112 | q.lock.Lock() 113 | defer q.lock.Unlock() 114 | q.users[name]-- 115 | if q.users[name] == 0 && qs.Cleared() { 116 | // Garbage collect unused queues. 117 | delete(q.users, name) 118 | delete(q.queues, name) 119 | } 120 | }() 121 | 122 | f(qs) 123 | } 124 | 125 | // Iterate calls f with every non-empty QueueState in q. 126 | func (q *QueueStateMux) Iterate(f func(string, *QueueState)) error { 127 | q.saveLock.RLock() 128 | defer q.saveLock.RUnlock() 129 | 130 | if q.finalError != nil { 131 | return q.finalError 132 | } 133 | 134 | q.lock.Lock() 135 | names := make([]string, 0, len(q.queues)) 136 | for name := range q.queues { 137 | names = append(names, name) 138 | } 139 | q.lock.Unlock() 140 | sort.Strings(names) 141 | for _, name := range names { 142 | q.get(name, func(qs *QueueState) { 143 | f(name, qs) 144 | }) 145 | } 146 | 147 | return nil 148 | } 149 | 150 | // Serialize writes the contents of the queue to a file, blocking all 151 | // operations on all queues to make sure cross-queue consistent state. 152 | // 153 | // If shutdown is true, then future operations after the state snapshot will 154 | // result in an ErrShuttingDown. 155 | func (q *QueueStateMux) Serialize(w io.Writer, shutdown bool) error { 156 | q.saveLock.Lock() 157 | if q.finalError != nil { 158 | return q.finalError 159 | } 160 | var states []ContextState 161 | for name, q := range q.queues { 162 | states = append(states, ContextState{ 163 | Name: name, 164 | Encoded: q.Encode(), 165 | }) 166 | } 167 | if shutdown { 168 | q.finalError = ErrShuttingDown 169 | } 170 | q.saveLock.Unlock() 171 | 172 | const context = "serialize queue state" 173 | 174 | resultWriter := zip.NewWriter(w) 175 | for i, state := range states { 176 | rw, err := resultWriter.Create(strconv.Itoa(i) + ".json") 177 | if err != nil { 178 | return errors.Wrap(err, context) 179 | } 180 | bufWriter := bufio.NewWriter(rw) 181 | if err := state.WriteJSON(bufWriter); err != nil { 182 | return errors.Wrap(err, context) 183 | } 184 | if err := bufWriter.Flush(); err != nil { 185 | return errors.Wrap(err, context) 186 | } 187 | } 188 | 189 | if err := resultWriter.Close(); err != nil { 190 | return errors.Wrap(err, context) 191 | } 192 | 193 | return nil 194 | } 195 | 196 | // QueueState maintains two queues of tasks: a pending queue and a running 197 | // queue. 198 | // 199 | // Tasks are added to the pending queue via Push(). When a task is returned 200 | // from Pop(), it is moved to the running queue and given an expiration time. 201 | // In general, Pop() first checks for tasks in the pending queue, and only 202 | // attempts to re-use an expired task from the running queue if necessary. 203 | // When Completed() is called for a task, it is removed from the running queue, 204 | // preventing it from ever being returned by Pop() again. 205 | // Tasks may be marked as completed at any time while they are in the running 206 | // queue, even if they are expired. 207 | type QueueState struct { 208 | lock sync.RWMutex 209 | pending *PendingQueue 210 | running *RunningQueue 211 | 212 | completionCounter int64 213 | lastModified time.Time 214 | rateTracker *RateTracker 215 | } 216 | 217 | // NewQueueState creates empty queues with the given task timeout. 218 | func NewQueueState(timeout time.Duration) *QueueState { 219 | return &QueueState{ 220 | pending: NewPendingQueue(), 221 | running: NewRunningQueue(timeout), 222 | lastModified: time.Now(), 223 | rateTracker: NewRateTracker(0), 224 | } 225 | } 226 | 227 | // DecodeQueueState decodes an object from QueueState.Encode() 228 | func DecodeQueueState(obj *EncodedQueueState) *QueueState { 229 | // Legacy tasks may have not stored a modtime, in which case 230 | // we update it to the time we load the checkpoint. 231 | var lastMod time.Time 232 | if obj.LastModified != nil { 233 | lastMod = *obj.LastModified 234 | } else { 235 | lastMod = time.Now() 236 | } 237 | 238 | return &QueueState{ 239 | pending: DecodePendingQueue(obj.Pending), 240 | running: DecodeRunningQueue(obj.Running), 241 | completionCounter: obj.Completed, 242 | lastModified: lastMod, 243 | rateTracker: DecodeRateTracker(obj.RateTracker), 244 | } 245 | } 246 | 247 | // Encode converts q into a JSON-serializable object. 248 | func (q *QueueState) Encode() *EncodedQueueState { 249 | q.lock.Lock() 250 | defer q.lock.Unlock() 251 | mt := q.lastModified 252 | return &EncodedQueueState{ 253 | Pending: q.pending.Encode(), 254 | Running: q.running.Encode(), 255 | Completed: q.completionCounter, 256 | LastModified: &mt, 257 | RateTracker: q.rateTracker.Encode(), 258 | } 259 | } 260 | 261 | // Push creates a task and returns the its new ID. 262 | // 263 | // If the specified maxSize is greater than 0, then the item will not be pushed 264 | // and false will be returned if the queue contains at least maxSize tasks. 265 | func (q *QueueState) Push(contents string, maxSize int) (string, bool) { 266 | q.lock.Lock() 267 | defer q.lock.Unlock() 268 | if maxSize > 0 && q.pending.Len()+q.running.Len() >= maxSize { 269 | return "", false 270 | } 271 | q.modified() 272 | return q.pending.AddTask(contents).ID, true 273 | } 274 | 275 | // PushBatch is like Push, except that it pushes multiple tasks at once. 276 | // 277 | // Either all or no tasks will be pushed depending on the maxSize and current 278 | // queue size. 279 | func (q *QueueState) PushBatch(contents []string, maxSize int) ([]string, bool) { 280 | q.lock.Lock() 281 | defer q.lock.Unlock() 282 | if maxSize > 0 && q.pending.Len()+q.running.Len()+len(contents) > maxSize { 283 | return nil, false 284 | } 285 | ids := make([]string, len(contents)) 286 | for i, x := range contents { 287 | ids[i] = q.pending.AddTask(x).ID 288 | } 289 | if len(contents) > 0 { 290 | q.modified() 291 | } 292 | return ids, true 293 | } 294 | 295 | // Pop gets a task from the queue, preferring the pending queue and dipping 296 | // into the expired tasks in the running queue only if necessary. 297 | func (q *QueueState) Pop(timeout *time.Duration) (*Task, *time.Time) { 298 | q.lock.Lock() 299 | defer q.lock.Unlock() 300 | nextPending := q.pending.PopTask() 301 | if nextPending != nil { 302 | q.modified() 303 | q.running.StartedTask(nextPending, timeout) 304 | return nextPending, nil 305 | } 306 | 307 | nextExpired, nextTry := q.running.PopExpired() 308 | if nextExpired != nil { 309 | q.modified() 310 | q.running.StartedTask(nextExpired, timeout) 311 | return nextExpired, nil 312 | } 313 | 314 | return nil, nextTry 315 | } 316 | 317 | // PopBatch atomically pops at most n tasks from the queue. 318 | // 319 | // If fewer than n tasks are returned, the second return value is the time that 320 | // the next running task will expire, or nil if no tasks were running before 321 | // PopBatch was called. 322 | func (q *QueueState) PopBatch(n int, timeout *time.Duration) ([]*Task, *time.Time) { 323 | q.lock.Lock() 324 | defer q.lock.Unlock() 325 | 326 | var tasks []*Task 327 | for len(tasks) < n { 328 | t := q.pending.PopTask() 329 | if t == nil { 330 | break 331 | } 332 | tasks = append(tasks, t) 333 | } 334 | var nextTry *time.Time 335 | for len(tasks) < n { 336 | var t *Task 337 | t, nextTry = q.running.PopExpired() 338 | if t == nil { 339 | break 340 | } 341 | tasks = append(tasks, t) 342 | } 343 | 344 | for _, t := range tasks { 345 | q.running.StartedTask(t, timeout) 346 | } 347 | if len(tasks) > 0 { 348 | q.modified() 349 | } 350 | 351 | return tasks, nextTry 352 | } 353 | 354 | // Peek gets the next available task to pop, if there is one. 355 | // 356 | // If no task is currently available, Peek returns the next task to expire and 357 | // the time when it will expire, or nil if no tasks are running. 358 | func (q *QueueState) Peek() (*Task, *Task, *time.Time) { 359 | q.lock.Lock() 360 | defer q.lock.Unlock() 361 | nextPending := q.pending.PeekTask() 362 | if nextPending != nil { 363 | return nextPending, nil, nil 364 | } 365 | return q.running.PeekExpired() 366 | } 367 | 368 | // Completed marks the identified task as complete, or returns false if no task 369 | // with the given ID was in the running queue. 370 | func (q *QueueState) Completed(id string) bool { 371 | q.lock.Lock() 372 | defer q.lock.Unlock() 373 | res := q.running.Completed(id) != nil 374 | if res { 375 | q.completionCounter += 1 376 | q.modified() 377 | q.rateTracker.Add(1) 378 | } 379 | return res 380 | } 381 | 382 | // Keepalive restarts the timeout period for the identified task, or returns 383 | // false if no task with the given ID was in the running queue. 384 | func (q *QueueState) Keepalive(id string, timeout *time.Duration) bool { 385 | q.lock.Lock() 386 | defer q.lock.Unlock() 387 | success := q.running.Keepalive(id, timeout) 388 | if success { 389 | q.modified() 390 | } 391 | return success 392 | } 393 | 394 | // Counts gets the current number of tasks in each state. 395 | func (q *QueueState) Counts(rateSeconds int, includeModtime, includeBytes bool) *QueueCounts { 396 | q.lock.RLock() 397 | defer q.lock.RUnlock() 398 | runningTotal := q.running.Len() 399 | runningExpired := q.running.NumExpired() 400 | var rate *float64 401 | if rateSeconds > 0 { 402 | rateSeconds = essentials.MinInt(rateSeconds, q.rateTracker.HistorySize()) 403 | r := float64(q.rateTracker.Count(rateSeconds)) / float64(rateSeconds) 404 | rate = &r 405 | } 406 | var modtime *int64 407 | if includeModtime { 408 | modtime = new(int64) 409 | *modtime = q.lastModified.UnixMilli() 410 | } 411 | var bytes *int64 412 | if includeBytes { 413 | bytes = new(int64) 414 | *bytes = q.pending.Bytes() + q.running.Bytes() 415 | } 416 | return &QueueCounts{ 417 | Pending: int64(q.pending.Len()), 418 | Running: int64(runningTotal - runningExpired), 419 | Expired: int64(runningExpired), 420 | Completed: q.completionCounter, 421 | LastModified: modtime, 422 | Rate: rate, 423 | Bytes: bytes, 424 | } 425 | } 426 | 427 | // Clear empties the queues and resets the completion counter. 428 | func (q *QueueState) Clear() { 429 | q.lock.Lock() 430 | defer q.lock.Unlock() 431 | q.pending.Clear() 432 | q.running.Clear() 433 | q.completionCounter = 0 434 | q.rateTracker.Reset() 435 | q.modified() 436 | } 437 | 438 | // Cleared returns true if the queue is effectively a fresh object, containing 439 | // no running tasks and zero completed tasks. 440 | func (q *QueueState) Cleared() bool { 441 | q.lock.RLock() 442 | defer q.lock.RUnlock() 443 | return q.pending.Len() == 0 && q.running.Len() == 0 && q.completionCounter == 0 444 | } 445 | 446 | // ExpireAll marks all tasks as expired, allowing them to be immediately popped 447 | // from the running queue. 448 | // 449 | // It does not move the tasks back to the pending queue. For this, call 450 | // QueueExpired(). 451 | func (q *QueueState) ExpireAll() int { 452 | q.lock.Lock() 453 | defer q.lock.Unlock() 454 | n := q.running.ExpireAll() 455 | if n > 0 { 456 | q.modified() 457 | } 458 | return n 459 | } 460 | 461 | // QueueExpired puts expired tasks from the running queue back into the pending 462 | // queue. 463 | func (q *QueueState) QueueExpired() int { 464 | q.lock.Lock() 465 | defer q.lock.Unlock() 466 | n := 0 467 | for { 468 | task, _ := q.running.PopExpired() 469 | if task == nil { 470 | break 471 | } 472 | n += 1 473 | q.pending.PushTask(task) 474 | } 475 | if n > 0 { 476 | q.modified() 477 | } 478 | return n 479 | } 480 | 481 | func (q *QueueState) modified() { 482 | q.lastModified = time.Now() 483 | } 484 | 485 | type PendingQueue struct { 486 | deque *TaskDeque 487 | curID int64 488 | } 489 | 490 | func NewPendingQueue() *PendingQueue { 491 | return &PendingQueue{deque: &TaskDeque{}} 492 | } 493 | 494 | // DecodePendingQueue decodes an object from PendingQueue.Encode(). 495 | func DecodePendingQueue(obj *EncodedPendingQueue) *PendingQueue { 496 | return &PendingQueue{ 497 | deque: DecodeTaskDeque(obj.Deque), 498 | curID: obj.CurID, 499 | } 500 | } 501 | 502 | // Encode converts p into a JSON-serializable object. 503 | func (p *PendingQueue) Encode() *EncodedPendingQueue { 504 | return &EncodedPendingQueue{ 505 | Deque: p.deque.Encode(), 506 | CurID: p.curID, 507 | } 508 | } 509 | 510 | // Bytes returns the number of bytes in all tasks in the queue. 511 | func (p *PendingQueue) Bytes() int64 { 512 | return p.deque.Bytes() 513 | } 514 | 515 | // AddTask creates a new task with the given contents and enqueues it. 516 | func (p *PendingQueue) AddTask(contents string) *Task { 517 | task := &Task{ 518 | Contents: contents, 519 | ID: strconv.FormatInt(p.curID, 16), 520 | } 521 | p.curID += 1 522 | p.deque.PushLast(task) 523 | return task 524 | } 525 | 526 | // PushTask re-enqueues an existing task. 527 | func (p *PendingQueue) PushTask(t *Task) { 528 | p.deque.PushLast(t) 529 | } 530 | 531 | // PopTask gets the next task (in FIFO order). 532 | func (p *PendingQueue) PopTask() *Task { 533 | return p.deque.PopFirst() 534 | } 535 | 536 | // PeekTask gets a copy of the next task. 537 | // 538 | // The copy only includes visible metadata. It will have no connection to the 539 | // queue or the original task. 540 | func (p *PendingQueue) PeekTask() *Task { 541 | t := p.deque.PeekFirst() 542 | if t == nil { 543 | return nil 544 | } 545 | return t.DisconnectedCopy() 546 | } 547 | 548 | // Len gets the number of queued tasks. 549 | func (p *PendingQueue) Len() int { 550 | return p.deque.Len() 551 | } 552 | 553 | // Clear deletes all of the pending tasks. 554 | func (p *PendingQueue) Clear() { 555 | p.deque = &TaskDeque{} 556 | } 557 | 558 | type RunningQueue struct { 559 | idToTask map[string]*Task 560 | deque *TaskDeque 561 | timeout time.Duration 562 | } 563 | 564 | func NewRunningQueue(timeout time.Duration) *RunningQueue { 565 | return &RunningQueue{ 566 | idToTask: map[string]*Task{}, 567 | deque: &TaskDeque{}, 568 | timeout: timeout, 569 | } 570 | } 571 | 572 | // DecodeRunningQueue decodes an object from RunningQueue.Encode(). 573 | func DecodeRunningQueue(obj *EncodedRunningQueue) *RunningQueue { 574 | deque := DecodeTaskDeque(obj.Deque) 575 | idToTask := map[string]*Task{} 576 | deque.Iterate(func(t *Task) { 577 | idToTask[t.ID] = t 578 | }) 579 | return &RunningQueue{ 580 | idToTask: idToTask, 581 | deque: deque, 582 | timeout: obj.Timeout, 583 | } 584 | } 585 | 586 | // Encode converts the queue into a JSON-serializable object. 587 | func (r *RunningQueue) Encode() *EncodedRunningQueue { 588 | return &EncodedRunningQueue{ 589 | Deque: r.deque.Encode(), 590 | Timeout: r.timeout, 591 | } 592 | } 593 | 594 | // Bytes returns the number of bytes in all tasks in the queue. 595 | func (r *RunningQueue) Bytes() int64 { 596 | return r.deque.Bytes() 597 | } 598 | 599 | // StartedTask adds the task to the queue and sets its timeout accordingly. 600 | func (r *RunningQueue) StartedTask(t *Task, timeout *time.Duration) { 601 | r.idToTask[t.ID] = t 602 | if timeout == nil { 603 | timeout = &r.timeout 604 | } 605 | t.expiration = time.Now().Add(*timeout) 606 | r.deque.PushByExpiration(t) 607 | } 608 | 609 | // PopExpired removes the first timed out task from the queue and returns it. 610 | // 611 | // If no tasks are timed out, the second return argument specifies the next 612 | // time when a task is set to expire (if there is one). 613 | func (r *RunningQueue) PopExpired() (*Task, *time.Time) { 614 | task := r.deque.PeekFirst() 615 | if task == nil { 616 | return nil, nil 617 | } 618 | now := time.Now() 619 | if task.expiration.After(now) { 620 | exp := task.expiration 621 | return nil, &exp 622 | } else { 623 | r.deque.Remove(task) 624 | delete(r.idToTask, task.ID) 625 | return task, nil 626 | } 627 | } 628 | 629 | // PeekExpired returns a copy of the first timed out task or the next task that 630 | // will expire in the queue. 631 | // 632 | // If no tasks are timed out, the second return value is the next task to 633 | // expire, and the third is the time when it will expire. 634 | // 635 | // If no tasks are enqueued (expired or not) all return values are nil. 636 | // 637 | // The returned tasks only include visible metadata. They will have no 638 | // connection to the queue or the original task. 639 | func (r *RunningQueue) PeekExpired() (*Task, *Task, *time.Time) { 640 | task := r.deque.PeekFirst() 641 | if task == nil { 642 | return nil, nil, nil 643 | } 644 | now := time.Now() 645 | if task.expiration.After(now) { 646 | exp := task.expiration 647 | return nil, task.DisconnectedCopy(), &exp 648 | } else { 649 | return task.DisconnectedCopy(), nil, nil 650 | } 651 | } 652 | 653 | // Completed removes a task from the queue. 654 | // 655 | // If the task is no longer in the queue, for example if it was removed with 656 | // PopExpired(), this returns nil. 657 | func (r *RunningQueue) Completed(id string) *Task { 658 | task, ok := r.idToTask[id] 659 | if !ok { 660 | return nil 661 | } 662 | r.deque.Remove(task) 663 | delete(r.idToTask, id) 664 | return task 665 | } 666 | 667 | // Keepalive restarts the timeout period for the identified task. 668 | // 669 | // Returns true if the task was found, or false otherwise. 670 | func (r *RunningQueue) Keepalive(id string, timeout *time.Duration) bool { 671 | task, ok := r.idToTask[id] 672 | if !ok { 673 | return false 674 | } 675 | r.deque.Remove(task) 676 | r.StartedTask(task, timeout) 677 | return true 678 | } 679 | 680 | // Len gets the number of tasks in the queue. 681 | func (r *RunningQueue) Len() int { 682 | return r.deque.Len() 683 | } 684 | 685 | // NumExpired gets the number of expired tasks. 686 | func (r *RunningQueue) NumExpired() int { 687 | now := time.Now() 688 | task := r.deque.first 689 | n := 0 690 | for task != nil && !task.expiration.After(now) { 691 | n++ 692 | task = task.queueNext 693 | } 694 | return n 695 | } 696 | 697 | // ExpireAll changes the timeout for all tasks to be before now. 698 | func (r *RunningQueue) ExpireAll() int { 699 | n := 0 700 | for _, task := range r.idToTask { 701 | n += 1 702 | task.expiration = time.Time{} 703 | } 704 | return n 705 | } 706 | 707 | // Clear deletes all of the running tasks. 708 | func (r *RunningQueue) Clear() { 709 | r.idToTask = map[string]*Task{} 710 | r.deque = &TaskDeque{} 711 | } 712 | 713 | type QueueCounts struct { 714 | Pending int64 `json:"pending"` 715 | Running int64 `json:"running"` 716 | Expired int64 `json:"expired"` 717 | Completed int64 `json:"completed"` 718 | LastModified *int64 `json:"modtime,omitempty"` 719 | Rate *float64 `json:"rate,omitempty"` 720 | Bytes *int64 `json:"bytes,omitempty"` 721 | } 722 | 723 | type ContextState struct { 724 | Name string 725 | Encoded *EncodedQueueState 726 | } 727 | 728 | func (c *ContextState) WriteJSON(w io.Writer) error { 729 | return WriteJSONObject(w, map[string]interface{}{ 730 | "Name": c.Name, 731 | "Encoded": c.Encoded, 732 | }) 733 | } 734 | 735 | type EncodedQueueState struct { 736 | Pending *EncodedPendingQueue 737 | Running *EncodedRunningQueue 738 | Completed int64 739 | LastModified *time.Time 740 | RateTracker *EncodedRateTracker 741 | } 742 | 743 | func (e *EncodedQueueState) WriteJSON(w io.Writer) error { 744 | t := e.LastModified 745 | return WriteJSONObject(w, map[string]interface{}{ 746 | "Pending": e.Pending, 747 | "Running": e.Running, 748 | "Completed": e.Completed, 749 | "LastModified": &t, 750 | "RateTracker": e.RateTracker, 751 | }) 752 | } 753 | 754 | type EncodedPendingQueue struct { 755 | Deque []EncodedTask 756 | CurID int64 757 | } 758 | 759 | func (e *EncodedPendingQueue) WriteJSON(w io.Writer) error { 760 | return WriteJSONObject(w, map[string]interface{}{ 761 | "Deque": EncodedTaskList(e.Deque), 762 | "CurID": e.CurID, 763 | }) 764 | } 765 | 766 | type EncodedRunningQueue struct { 767 | Deque []EncodedTask 768 | Timeout time.Duration 769 | } 770 | 771 | func (e *EncodedRunningQueue) WriteJSON(w io.Writer) error { 772 | return WriteJSONObject(w, map[string]interface{}{ 773 | "Deque": EncodedTaskList(e.Deque), 774 | "Timeout": e.Timeout, 775 | }) 776 | } 777 | -------------------------------------------------------------------------------- /tasq-server/rate_tracker.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // Default history time used for RateTracker. 8 | const DefaultRateTrackerBins = 128 9 | 10 | // A RateTracker keeps a sliding window of event counts over the last 11 | // N seconds. 12 | type RateTracker struct { 13 | firstBinTime int64 14 | bins []int64 15 | } 16 | 17 | // NewRateTracker creates a RateTracker which keeps event counts up to 18 | // historySize seconds in the past. 19 | // 20 | // If historySize is 0, then DefaultRateTrackerBins is used. 21 | func NewRateTracker(historySize int) *RateTracker { 22 | if historySize == 0 { 23 | historySize = DefaultRateTrackerBins 24 | } 25 | return &RateTracker{ 26 | bins: make([]int64, historySize), 27 | } 28 | } 29 | 30 | // DecodeRateTracker loads an encoded RateTracker. 31 | // If the state is empty, a new rate tracker with DefaultRateTrackerBins is 32 | // created. 33 | func DecodeRateTracker(state *EncodedRateTracker) *RateTracker { 34 | if state == nil || len(state.Bins) == 0 { 35 | return NewRateTracker(DefaultRateTrackerBins) 36 | } 37 | return &RateTracker{ 38 | firstBinTime: state.FirstBinTime, 39 | bins: state.Bins, 40 | } 41 | } 42 | 43 | // Reset zeros out the counters. 44 | func (r *RateTracker) Reset() { 45 | for i := range r.bins { 46 | r.bins[i] = 0 47 | } 48 | } 49 | 50 | // HistorySize returns the number of time bins. 51 | func (r *RateTracker) HistorySize() int { 52 | return len(r.bins) 53 | } 54 | 55 | // Add adds the count n to the current time bin. 56 | func (r *RateTracker) Add(n int64) { 57 | r.AddAt(time.Now().Unix(), n) 58 | } 59 | 60 | // AddAt is like Add, but allows the caller to specify the current time. 61 | func (r *RateTracker) AddAt(curTime, n int64) { 62 | r.truncateAndShift(curTime) 63 | r.bins[len(r.bins)-1] += n 64 | } 65 | 66 | // Count retrieves the count over the last t seconds. 67 | // The t argument must be at most the history size passed to NewRateTracker. 68 | func (r *RateTracker) Count(t int) int64 { 69 | return r.CountAt(time.Now().Unix(), t) 70 | } 71 | 72 | // CountAt is like Count, but allows the caller to specify the current time. 73 | func (r *RateTracker) CountAt(curTime int64, t int) int64 { 74 | if t > len(r.bins) { 75 | panic("too many seconds requested") 76 | } 77 | r.truncateAndShift(curTime) 78 | var res int64 79 | for i := len(r.bins) - 1; i >= len(r.bins)-t; i-- { 80 | res += r.bins[i] 81 | } 82 | return res 83 | } 84 | 85 | func (r *RateTracker) Encode() *EncodedRateTracker { 86 | return &EncodedRateTracker{ 87 | FirstBinTime: r.firstBinTime, 88 | Bins: append([]int64{}, r.bins...), 89 | } 90 | } 91 | 92 | func (r *RateTracker) truncateAndShift(curTime int64) { 93 | lastBinTime := r.firstBinTime + int64(len(r.bins)) - 1 94 | 95 | if curTime < r.firstBinTime || curTime >= lastBinTime+int64(len(r.bins)) { 96 | r.firstBinTime = curTime - (int64(len(r.bins)) - 1) 97 | for i := range r.bins { 98 | r.bins[i] = 0 99 | } 100 | } else if curTime < lastBinTime { 101 | // The clock has likely changed a tiny bit, so the last bin is in 102 | // the future. This rarely happens, but we lose history when it does. 103 | backtrack := lastBinTime - curTime 104 | r.firstBinTime -= backtrack 105 | copy(r.bins[backtrack:], r.bins[:]) 106 | for i := 0; i < int(backtrack); i++ { 107 | r.bins[i] = 0 108 | } 109 | } else if curTime > lastBinTime { 110 | forward := curTime - lastBinTime 111 | r.firstBinTime += forward 112 | copy(r.bins[:], r.bins[forward:]) 113 | for i := len(r.bins) - int(forward); i < len(r.bins); i++ { 114 | r.bins[i] = 0 115 | } 116 | } 117 | } 118 | 119 | type EncodedRateTracker struct { 120 | FirstBinTime int64 121 | Bins []int64 122 | } 123 | -------------------------------------------------------------------------------- /tasq-server/rate_tracker_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestRateTracker(t *testing.T) { 8 | rt := NewRateTracker(5) 9 | rt.AddAt(30, 2) 10 | rt.AddAt(31, 3) 11 | rt.AddAt(32, 4) 12 | rt.AddAt(32, 1) 13 | rt.AddAt(33, 1) 14 | if count := rt.CountAt(33, 1); count != 1 { 15 | t.Fatalf("bad count: %d", count) 16 | } 17 | if count := rt.CountAt(33, 2); count != 6 { 18 | t.Fatalf("bad count: %d", count) 19 | } 20 | if count := rt.CountAt(33, 5); count != 11 { 21 | t.Fatalf("bad count: %d", count) 22 | } 23 | rt.AddAt(34, 3) 24 | if count := rt.CountAt(34, 2); count != 4 { 25 | t.Fatalf("bad count: %d", count) 26 | } 27 | rt.AddAt(36, 7) 28 | if count := rt.CountAt(36, 1); count != 7 { 29 | t.Fatalf("bad count: %d", count) 30 | } 31 | if count := rt.CountAt(36, 2); count != 7 { 32 | t.Fatalf("bad count: %d", count) 33 | } 34 | if count := rt.CountAt(36, 3); count != 10 { 35 | t.Fatalf("bad count: %d", count) 36 | } 37 | if count := rt.CountAt(36, 5); count != 16 { 38 | t.Fatalf("bad count: %d", count) 39 | } 40 | rt.AddAt(35, 5) 41 | if count := rt.CountAt(35, 1); count != 5 { 42 | t.Fatalf("bad count: %d", count) 43 | } 44 | if count := rt.CountAt(35, 2); count != 8 { 45 | t.Fatalf("bad count: %d", count) 46 | } 47 | if count := rt.CountAt(35, 5); count != 14 { 48 | t.Fatalf("bad count: %d", count) 49 | } 50 | rt.AddAt(40, 10) 51 | if count := rt.CountAt(40, 5); count != 10 { 52 | t.Fatalf("bad count: %d", count) 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /tasq-server/task.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import "time" 4 | 5 | type Task struct { 6 | ID string `json:"id"` 7 | Contents string `json:"contents"` 8 | 9 | // For in-progress tasks. 10 | expiration time.Time 11 | 12 | queuePrev *Task 13 | queueNext *Task 14 | } 15 | 16 | func (t *Task) DisconnectedCopy() *Task { 17 | return &Task{ID: t.ID, Contents: t.Contents} 18 | } 19 | 20 | type TaskDeque struct { 21 | first *Task 22 | last *Task 23 | count int 24 | bytes int64 25 | } 26 | 27 | // DecodeTaskDeque inverts TaskDeque.Encode(), converting a serializable 28 | // object back into a linked list deque. 29 | func DecodeTaskDeque(obj []EncodedTask) *TaskDeque { 30 | res := &TaskDeque{count: len(obj)} 31 | for i, et := range obj { 32 | task := &Task{ID: et.ID, Contents: et.Contents, expiration: et.Expiration} 33 | if i == 0 { 34 | res.first = task 35 | res.last = task 36 | } else { 37 | res.last.queueNext = task 38 | task.queuePrev = res.last 39 | res.last = task 40 | } 41 | res.bytes += int64(len(et.Contents)) 42 | } 43 | return res 44 | } 45 | 46 | // Encode generates a JSON-serializable object for the task sequence. 47 | // This can be reversed by DecodeTaskDeque. 48 | func (t *TaskDeque) Encode() []EncodedTask { 49 | objs := make([]EncodedTask, 0, t.count) 50 | t.Iterate(func(obj *Task) { 51 | objs = append(objs, EncodedTask{ 52 | ID: obj.ID, 53 | Contents: obj.Contents, 54 | Expiration: obj.expiration, 55 | }) 56 | }) 57 | return objs 58 | } 59 | 60 | func (t *TaskDeque) Len() int { 61 | return t.count 62 | } 63 | 64 | func (t *TaskDeque) Bytes() int64 { 65 | return t.bytes 66 | } 67 | 68 | func (t *TaskDeque) PushLast(task *Task) { 69 | t.count += 1 70 | t.bytes += int64(len(task.Contents)) 71 | if t.last == nil { 72 | t.first = task 73 | t.last = task 74 | task.queuePrev = nil 75 | task.queueNext = nil 76 | } else { 77 | t.last.queueNext = task 78 | task.queuePrev = t.last 79 | task.queueNext = nil 80 | t.last = task 81 | } 82 | } 83 | 84 | func (t *TaskDeque) PushFirst(task *Task) { 85 | t.count += 1 86 | t.bytes += int64(len(task.Contents)) 87 | if t.first == nil { 88 | t.first = task 89 | t.last = task 90 | task.queuePrev = nil 91 | task.queueNext = nil 92 | } else { 93 | t.first.queuePrev = task 94 | task.queueNext = t.first 95 | task.queuePrev = nil 96 | t.first = task 97 | } 98 | } 99 | 100 | func (t *TaskDeque) PushByExpiration(task *Task) { 101 | prev := t.last 102 | for prev != nil && prev.expiration.After(task.expiration) { 103 | prev = prev.queuePrev 104 | } 105 | if prev == nil { 106 | t.PushFirst(task) 107 | } else if prev.queueNext == nil { 108 | t.PushLast(task) 109 | } else { 110 | t.bytes += int64(len(task.Contents)) 111 | t.count += 1 112 | next := prev.queueNext 113 | prev.queueNext = task 114 | next.queuePrev = task 115 | task.queuePrev = prev 116 | task.queueNext = next 117 | } 118 | } 119 | 120 | func (t *TaskDeque) PopLast() *Task { 121 | res := t.last 122 | if res != nil { 123 | t.Remove(res) 124 | } 125 | return res 126 | } 127 | 128 | func (t *TaskDeque) PopFirst() *Task { 129 | res := t.first 130 | if res != nil { 131 | t.Remove(res) 132 | } 133 | return res 134 | } 135 | 136 | func (t *TaskDeque) PeekFirst() *Task { 137 | return t.first 138 | } 139 | 140 | func (t *TaskDeque) Remove(task *Task) { 141 | if task.queuePrev == nil { 142 | if t.first != task { 143 | panic("task not in deque") 144 | } 145 | t.first = task.queueNext 146 | task.queueNext = nil 147 | if t.first != nil { 148 | t.first.queuePrev = nil 149 | } else { 150 | t.last = nil 151 | } 152 | } else if task.queueNext == nil { 153 | if t.last != task { 154 | panic("task not in queue") 155 | } 156 | t.last = task.queuePrev 157 | task.queuePrev = nil 158 | if t.last != nil { 159 | t.last.queueNext = nil 160 | } else { 161 | t.first = nil 162 | } 163 | } else { 164 | task.queueNext.queuePrev = task.queuePrev 165 | task.queuePrev.queueNext = task.queueNext 166 | task.queueNext = nil 167 | task.queuePrev = nil 168 | } 169 | t.count -= 1 170 | t.bytes -= int64(len(task.Contents)) 171 | if task.queueNext != nil || task.queuePrev != nil { 172 | panic("pointer unexpectedly preserved") 173 | } 174 | } 175 | 176 | func (t *TaskDeque) Iterate(f func(t *Task)) { 177 | obj := t.first 178 | for obj != nil { 179 | f(obj) 180 | obj = obj.queueNext 181 | } 182 | } 183 | 184 | type EncodedTask struct { 185 | ID string 186 | Contents string 187 | Expiration time.Time 188 | } 189 | -------------------------------------------------------------------------------- /tasq-transfer/main.go: -------------------------------------------------------------------------------- 1 | // Command tasq-transfer moves tasks from one tasq server to another. 2 | // 3 | // Regardless of program interruption or network errors, no tasks will be lost. 4 | // In particular, a crash or network failure during the transfer may result in 5 | // some tasks being duplicated between the source and destination servers, but 6 | // no tasks will be removed from the source before being added to the 7 | // destination. 8 | package main 9 | 10 | import ( 11 | "flag" 12 | "log" 13 | "time" 14 | 15 | "github.com/unixpickle/essentials" 16 | "github.com/unixpickle/tasq" 17 | ) 18 | 19 | func main() { 20 | var sourceHost string 21 | var sourceContext string 22 | var sourceUsername string 23 | var sourcePassword string 24 | var destHost string 25 | var destContext string 26 | var destUsername string 27 | var destPassword string 28 | var numTasks int 29 | var bufferSize int 30 | var waitRunning bool 31 | flag.StringVar(&sourceHost, "source", "", "source server URL") 32 | flag.StringVar(&sourceContext, "source-context", "", "source context") 33 | flag.StringVar(&sourceUsername, "source-username", "", "source basic auth username") 34 | flag.StringVar(&sourcePassword, "source-password", "", "source basic auth password") 35 | flag.StringVar(&destHost, "dest", "", "destination server URL") 36 | flag.StringVar(&destContext, "dest-context", "", "destination context") 37 | flag.StringVar(&destUsername, "dest-username", "", "destination basic auth username") 38 | flag.StringVar(&destPassword, "dest-password", "", "destination basic auth password") 39 | flag.IntVar(&numTasks, "count", -1, "number of tasks to transfer") 40 | flag.IntVar(&bufferSize, "buffer-size", 4096, "task buffer size") 41 | flag.BoolVar(&waitRunning, "wait-running", false, 42 | "attempt to transfer in-progress tasks once they expire") 43 | flag.Parse() 44 | 45 | if sourceHost == "" || destHost == "" { 46 | essentials.Die("Must provide -source and -dest. See -help.") 47 | } 48 | 49 | sourceClient, err := tasq.NewClient(sourceHost, sourceContext, sourceUsername, sourcePassword) 50 | essentials.Must(err) 51 | 52 | destClient, err := tasq.NewClient(destHost, destContext, destUsername, destPassword) 53 | essentials.Must(err) 54 | 55 | completed := 0 56 | for numTasks == -1 || completed < numTasks { 57 | bs := bufferSize 58 | if numTasks != -1 && bs > numTasks-completed { 59 | bs = numTasks - completed 60 | } 61 | tasks, retry, err := sourceClient.PopBatch(bs) 62 | if err != nil { 63 | log.Fatalln("ERROR popping batch:", err) 64 | } 65 | if len(tasks) == 0 && retry == nil { 66 | log.Println("Source queue has been exhausted.") 67 | break 68 | } else if len(tasks) == 0 { 69 | if waitRunning { 70 | log.Printf("Waiting %f seconds for next timeout...", *retry) 71 | time.Sleep(time.Duration(float64(time.Second) * *retry)) 72 | } else { 73 | log.Printf("Done all immediately available tasks (wait time %f).", *retry) 74 | break 75 | } 76 | } else { 77 | var ids, contents []string 78 | for _, t := range tasks { 79 | ids = append(ids, t.ID) 80 | contents = append(contents, t.Contents) 81 | } 82 | if _, err := destClient.PushBatch(contents); err != nil { 83 | log.Fatalln("ERROR pushing batch:", err) 84 | } 85 | if err := sourceClient.CompletedBatch(ids); err != nil { 86 | log.Fatalln("ERROR marking batch as completed:", err) 87 | } 88 | completed += len(tasks) 89 | log.Printf("Current status: transferred a total of %d tasks", completed) 90 | } 91 | } 92 | } 93 | --------------------------------------------------------------------------------