feat: Initial commit
This commit is contained in:
mrhid6
349
handler.go
Normal file
349
handler.go
Normal file
@@ -0,0 +1,349 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"taskqueue/models"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
// ---------- Handlers ----------
|
||||
|
||||
func (s *Server) handleEnqueue(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
httpError(w, http.StatusMethodNotAllowed, "POST required")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
ApplicationId string `json:"applicationId"`
|
||||
Payload map[string]interface{} `json:"payload"`
|
||||
Priority int `json:"priority"`
|
||||
DelaySec int `json:"delaySec"` // optional delay until available
|
||||
MaxAttempts int `json:"maxAttempts"` // optional
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
availableAt := now
|
||||
if req.DelaySec > 0 {
|
||||
availableAt = now.Add(time.Duration(req.DelaySec) * time.Second)
|
||||
}
|
||||
if req.MaxAttempts <= 0 {
|
||||
req.MaxAttempts = 5
|
||||
}
|
||||
|
||||
task := models.Task{
|
||||
ApplicationId: req.ApplicationId,
|
||||
Payload: req.Payload,
|
||||
Priority: req.Priority,
|
||||
AvailableAt: availableAt,
|
||||
Status: StatusPending,
|
||||
Attempts: 0,
|
||||
MaxAttempts: req.MaxAttempts,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
res, err := s.col.InsertOne(ctx, task)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, "insert: "+err.Error())
|
||||
return
|
||||
}
|
||||
task.ID = res.InsertedID.(primitive.ObjectID)
|
||||
writeJSON(w, http.StatusCreated, task)
|
||||
}
|
||||
|
||||
func (s *Server) handlePop(w http.ResponseWriter, r *http.Request) {
|
||||
// Atomically find a task matching filters and claim it (lease)
|
||||
if r.Method != http.MethodPost {
|
||||
httpError(w, http.StatusMethodNotAllowed, "POST required")
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
WorkerID string `json:"workerId"` // required - identity of popper
|
||||
LeaseSeconds int `json:"leaseSeconds"` // optional, defaults to 60
|
||||
Filters map[string]interface{} `json:"filters"` // e.g. {"type":"email"}
|
||||
MinPriority *int `json:"minPriority"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.WorkerID == "" {
|
||||
httpError(w, http.StatusBadRequest, "workerId required")
|
||||
return
|
||||
}
|
||||
leaseSeconds := req.LeaseSeconds
|
||||
if leaseSeconds <= 0 {
|
||||
leaseSeconds = 60
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
leaseUntil := now.Add(time.Duration(leaseSeconds) * time.Second)
|
||||
|
||||
// Build filter:
|
||||
// - task is pending and availableAt <= now
|
||||
// OR - task is leased but leasedUntil <= now (expired)
|
||||
filter := bson.A{}
|
||||
|
||||
// user filters (e.g. type)
|
||||
userFilter := bson.D{}
|
||||
for k, v := range req.Filters {
|
||||
userFilter = append(userFilter, bson.E{Key: k, Value: v})
|
||||
}
|
||||
if req.MinPriority != nil {
|
||||
userFilter = append(userFilter, bson.E{Key: "priority", Value: bson.D{{"$gte", *req.MinPriority}}})
|
||||
}
|
||||
|
||||
criteria := bson.D{
|
||||
{"$and", bson.A{
|
||||
userFilter,
|
||||
bson.D{{"availableAt", bson.D{{"$lte", now}}}},
|
||||
bson.D{{"$or", bson.A{
|
||||
bson.D{{"status", StatusPending}},
|
||||
bson.D{{"status", StatusLeased}, {"leasedUntil", bson.D{{"$lte", now}}}}, // expired
|
||||
}}},
|
||||
}},
|
||||
}
|
||||
|
||||
filter = append(filter, criteria)
|
||||
|
||||
findFilter := filter[0]
|
||||
|
||||
// Sort by priority desc, createdAt asc
|
||||
sort := bson.D{{"priority", -1}, {"createdAt", 1}}
|
||||
|
||||
update := bson.D{
|
||||
{"$set", bson.D{
|
||||
{"status", StatusLeased},
|
||||
{"leaseOwner", req.WorkerID},
|
||||
{"leasedUntil", leaseUntil},
|
||||
{"updatedAt", now},
|
||||
}},
|
||||
{"$inc", bson.D{{"attempts", 1}}},
|
||||
}
|
||||
|
||||
opts := options.FindOneAndUpdate().SetSort(sort).SetReturnDocument(options.After)
|
||||
ctx := r.Context()
|
||||
|
||||
var task models.Task
|
||||
err := s.col.FindOneAndUpdate(ctx, findFilter.(bson.D), update, opts).Decode(&task)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||||
// nothing to pop
|
||||
writeJSON(w, http.StatusNoContent, nil)
|
||||
return
|
||||
}
|
||||
httpError(w, http.StatusInternalServerError, "pop error: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// If attempts > maxAttempts, mark done with error note instead of returning
|
||||
if task.MaxAttempts > 0 && task.Attempts > task.MaxAttempts {
|
||||
_, _ = s.col.UpdateByID(ctx, task.ID, bson.D{
|
||||
{"$set", bson.D{
|
||||
{"status", StatusDone},
|
||||
{"error", "max attempts exceeded"},
|
||||
{"updatedAt", time.Now().UTC()},
|
||||
}},
|
||||
})
|
||||
writeJSON(w, http.StatusNoContent, nil)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, task)
|
||||
}
|
||||
|
||||
func (s *Server) handleComplete(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
httpError(w, http.StatusMethodNotAllowed, "POST required")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
TaskID string `json:"taskId"`
|
||||
WorkerID string `json:"workerId"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.TaskID == "" || req.WorkerID == "" {
|
||||
httpError(w, http.StatusBadRequest, "taskId and workerId required")
|
||||
return
|
||||
}
|
||||
oid, err := primitive.ObjectIDFromHex(req.TaskID)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid taskId")
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
now := time.Now().UTC()
|
||||
// Only allow completion if worker matches leaseOwner (or lease expired)
|
||||
filter := bson.D{
|
||||
{"_id", oid},
|
||||
{"$or", bson.A{
|
||||
bson.D{{"leaseOwner", req.WorkerID}},
|
||||
bson.D{{"leasedUntil", bson.D{{"$lte", now}}}}, // lease expired -> allow completion as cleanup
|
||||
}},
|
||||
}
|
||||
update := bson.D{
|
||||
{"$set", bson.D{
|
||||
{"status", StatusDone},
|
||||
{"updatedAt", now},
|
||||
}},
|
||||
}
|
||||
res, err := s.col.UpdateOne(ctx, filter, update)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, "complete update: "+err.Error())
|
||||
return
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
httpError(w, http.StatusNotFound, "task not found or worker mismatch")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"result": "ok"})
|
||||
}
|
||||
|
||||
func (s *Server) handleExtend(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
httpError(w, http.StatusMethodNotAllowed, "POST required")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
TaskID string `json:"taskId"`
|
||||
WorkerID string `json:"workerId"`
|
||||
ExtraSecond int `json:"extraSecond"` // how many more seconds to add
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.TaskID == "" || req.WorkerID == "" || req.ExtraSecond <= 0 {
|
||||
httpError(w, http.StatusBadRequest, "taskId, workerId and extraSecond required")
|
||||
return
|
||||
}
|
||||
oid, err := primitive.ObjectIDFromHex(req.TaskID)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid taskId")
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
now := time.Now().UTC()
|
||||
|
||||
filter := bson.D{
|
||||
{"_id", oid},
|
||||
{"leaseOwner", req.WorkerID},
|
||||
{"leasedUntil", bson.D{{"$gt", now}}}, // lease must still be valid
|
||||
}
|
||||
|
||||
// do proper update: set leasedUntil to old + extra
|
||||
var task models.Task
|
||||
err = s.col.FindOne(ctx, filter).Decode(&task)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, "find task: "+err.Error())
|
||||
return
|
||||
}
|
||||
if task.LeaseOwner != req.WorkerID {
|
||||
httpError(w, http.StatusForbidden, "worker does not own lease")
|
||||
return
|
||||
}
|
||||
newUntil := task.LeasedUntil.Add(time.Duration(req.ExtraSecond) * time.Second)
|
||||
_, err = s.col.UpdateOne(ctx, filter, bson.D{
|
||||
{"$set", bson.D{
|
||||
{"leasedUntil", newUntil},
|
||||
{"updatedAt", now},
|
||||
}},
|
||||
})
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, "extend update: "+err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"result": "ok", "newLeasedUntil": newUntil.Format(time.RFC3339)})
|
||||
}
|
||||
|
||||
func (s *Server) handleRequeue(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
httpError(w, http.StatusMethodNotAllowed, "POST required")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
TaskID string `json:"taskId"`
|
||||
WorkerID string `json:"workerId"`
|
||||
DelaySec int `json:"delaySec"` // optional
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.TaskID == "" || req.WorkerID == "" {
|
||||
httpError(w, http.StatusBadRequest, "taskId and workerId required")
|
||||
return
|
||||
}
|
||||
oid, err := primitive.ObjectIDFromHex(req.TaskID)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid taskId")
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
availableAt := now
|
||||
if req.DelaySec > 0 {
|
||||
availableAt = now.Add(time.Duration(req.DelaySec) * time.Second)
|
||||
}
|
||||
ctx := r.Context()
|
||||
|
||||
// Only allow requeue if worker owns lease OR lease expired.
|
||||
// We'll increment attempts only when popped, so here we don't change attempts.
|
||||
filter := bson.D{
|
||||
{"_id", oid},
|
||||
{"$or", bson.A{
|
||||
bson.D{{"leaseOwner", req.WorkerID}},
|
||||
bson.D{{"leasedUntil", bson.D{{"$lte", now}}}},
|
||||
}},
|
||||
}
|
||||
update := bson.D{
|
||||
{"$set", bson.D{
|
||||
{"status", StatusPending},
|
||||
{"availableAt", availableAt},
|
||||
{"leaseOwner", ""},
|
||||
{"leasedUntil", time.Time{}},
|
||||
{"updatedAt", now},
|
||||
{"error", req.Error},
|
||||
}},
|
||||
}
|
||||
res, err := s.col.UpdateOne(ctx, filter, update)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, "requeue update: "+err.Error())
|
||||
return
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
httpError(w, http.StatusNotFound, "task not found or worker mismatch")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"result": "ok"})
|
||||
}
|
||||
|
||||
// ---------- Helpers ----------
|
||||
func httpError(w http.ResponseWriter, status int, msg string) {
|
||||
w.WriteHeader(status)
|
||||
writeJSON(w, status, map[string]string{"error": msg})
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, v interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
Reference in New Issue
Block a user