Merge pull request #1559 from niaow/id-alloc-desync-structured-error

[CORE-386] Change ID allocation to return a structured error on offset desync

api_test.go

@@ -17,6 +17,7 @@ package pilosa_test
 import (
 	"context"
 	"crypto/rand"
+	"errors"
 	"fmt"
 	"math"
 	"reflect"
@@ -699,38 +700,116 @@ func TestAPI_IDAlloc(t *testing.T) {
 
 	primary := c.GetPrimary().API
 
-	key := pilosa.IDAllocKey{
-		Index: "index",
-		Key:   "key",
-	}
-	var session [32]byte
-	_, err := rand.Read(session[:])
-	if err != nil {
-		t.Fatalf("obtaining random bytes: %v", err)
-	}
+	t.Run("Normal", func(t *testing.T) {
+		key := pilosa.IDAllocKey{
+			Index: "normal",
+			Key:   "key",
+		}
+		var session [32]byte
+		_, err := rand.Read(session[:])
+		if err != nil {
+			t.Fatalf("obtaining random bytes: %v", err)
+		}
 
-	const toReserve = 2
+		const toReserve = 2
 
-	ids, err := primary.ReserveIDs(key, session, ^uint64(0), toReserve)
-	if err != nil {
-		t.Fatalf("reserving IDs: %v", err)
-	}
+		ids, err := primary.ReserveIDs(key, session, ^uint64(0), toReserve)
+		if err != nil {
+			t.Fatalf("reserving IDs: %v", err)
+		}
 
-	var numIds uint64
-	for _, idr := range ids {
-		numIds += (idr.Last - idr.First) + 1
-	}
-	if numIds != toReserve {
-		t.Errorf("expected %d ids but got %d: %v", toReserve, numIds, ids)
-	}
+		var numIds uint64
+		for _, idr := range ids {
+			numIds += (idr.Last - idr.First) + 1
+		}
+		if numIds != toReserve {
+			t.Errorf("expected %d ids but got %d: %v", toReserve, numIds, ids)
+		}
 
-	err = primary.CommitIDs(key, session, numIds)
-	if err != nil {
-		t.Fatalf("committing IDs: %v", err)
-	}
+		err = primary.CommitIDs(key, session, numIds)
+		if err != nil {
+			t.Fatalf("committing IDs: %v", err)
+		}
 
-	err = primary.ResetIDAlloc(key.Index)
-	if err != nil {
-		t.Fatalf("resetting ID alloc: %v", err)
-	}
+		err = primary.ResetIDAlloc(key.Index)
+		if err != nil {
+			t.Fatalf("resetting ID alloc: %v", err)
+		}
+	})
+	t.Run("Offset", func(t *testing.T) {
+		key := pilosa.IDAllocKey{
+			Index: "offset",
+			Key:   "key",
+		}
+		var session [32]byte
+		_, err := rand.Read(session[:])
+		if err != nil {
+			t.Fatalf("obtaining random bytes: %v", err)
+		}
+
+		ids, err := primary.ReserveIDs(key, session, 0, 2)
+		if err != nil {
+			t.Fatalf("reserving IDs: %v", err)
+		}
+
+		{
+			var numIds uint64
+			for _, idr := range ids {
+				numIds += (idr.Last - idr.First) + 1
+			}
+			if numIds != 2 {
+				t.Errorf("expected %d ids but got %d: %v", 2, numIds, ids)
+			}
+		}
+
+		_, err = rand.Read(session[:])
+		if err != nil {
+			t.Fatalf("obtaining random bytes: %v", err)
+		}
+		ids2, err := primary.ReserveIDs(key, session, 1, 2)
+		if err != nil {
+			t.Fatalf("reserving IDs with partially increased offset: %v", err)
+		}
+
+		var numIds uint64
+		for _, idr := range ids2 {
+			numIds += (idr.Last - idr.First) + 1
+		}
+		if numIds != 2 {
+			t.Errorf("expected %d ids but got %d: %v", 2, numIds, ids2)
+		}
+
+		if prevEnd, newStart := ids[len(ids)-1].Last, ids2[0].First; prevEnd != newStart {
+			t.Errorf("expected reuse of last ID (%d), but started with %d", prevEnd, newStart)
+		}
+
+		err = primary.CommitIDs(key, session, numIds)
+		if err != nil {
+			t.Errorf("committing IDs: %v", err)
+		}
+
+		_, err = rand.Read(session[:])
+		if err != nil {
+			t.Fatalf("obtaining random bytes: %v", err)
+		}
+		ids3, err := primary.ReserveIDs(key, session, 0, 2)
+		var esync pilosa.ErrIDOffsetDesync
+		if errors.As(err, &esync) {
+			if esync.Requested != 0 {
+				t.Errorf("incorrect requested offset in error: provided %d but got %d", 0, esync.Requested)
+			}
+			if esync.Base != 3 {
+				t.Errorf("incorrect base offset: expected %d but got %d", 3, esync.Base)
+			}
+		} else if err == nil {
+			t.Errorf("successfully re-reserved at a committed offset: %v", ids3)
+		} else {
+			t.Fatalf("unexpected error when reserving committed IDs: %v", err)
+		}
+
+		err = primary.ResetIDAlloc(key.Index)
+		if err != nil {
+			t.Fatalf("resetting ID alloc: %v", err)
+		}
+	})
 }

http/handler.go

@@ -2921,10 +2921,27 @@ func (h *Handler) handleReserveIDs(w http.ResponseWriter, r *http.Request) {
 
 	ids, err := h.api.ReserveIDs(req.Key, req.Session, req.Offset, req.Count)
 	if err != nil {
+		var esync pilosa.ErrIDOffsetDesync
+		if errors.As(err, &esync) {
+			w.Header().Add("Content-Type", "application/json")
+			w.WriteHeader(http.StatusConflict)
+			err = json.NewEncoder(w).Encode(struct {
+				pilosa.ErrIDOffsetDesync
+				Err string `json:"error"`
+			}{
+				ErrIDOffsetDesync: esync,
+				Err:               err.Error(),
+			})
+			if err != nil {
+				h.logger.Debugf("failed to send desync error: %v", err)
+			}
+			return
+		}
 		http.Error(w, fmt.Sprintf("reserving IDs: %v", err.Error()), http.StatusBadRequest)
 		return
 	}
 
+	w.Header().Add("Content-Type", "application/json")
 	err = json.NewEncoder(w).Encode(ids)
 	if err != nil {
 		http.Error(w, "encoding result", http.StatusBadRequest)

idalloc.go

@@ -16,6 +16,7 @@ package pilosa
 
 import (
 	"encoding/binary"
+	"fmt"
 	"math/bits"
 	"sort"
 	"time"
@@ -71,6 +72,21 @@ func (ida *idAllocator) Close() error {
 	return ida.db.Close()
 }
 
+// ErrIDOffsetDesync is an error generated when attempting to reserve IDs at a committed offset.
+// This will typically happen when kafka partitions are moved between kafka ingesters - there may be a brief period in which 2 ingesters are processing the same messages at the same time.
+// The ingester can resolve this by ignoring messages under base.
+type ErrIDOffsetDesync struct {
+	// Requested is the offset that the client attempted to reserve.
+	Requested uint64 `json:"requested"`
+
+	// Base is the next uncommitted offset for which IDs may be reserved.
+	Base uint64 `json:"base"`
+}
+
+func (err ErrIDOffsetDesync) Error() string {
+	return fmt.Sprintf("attempted to reserve IDs at committed offset %d (base offset: %d)", err.Requested, err.Base)
+}
+
 func (ida *idAllocator) reserve(key IDAllocKey, session [32]byte, offset, count uint64) ([]IDRange, error) {
 	if session == [32]byte{} {
 		return nil, errors.New("detected a broken session key")
@@ -94,9 +110,12 @@ func (ida *idAllocator) reserve(key IDAllocKey, session [32]byte, offset, count
 		if offset != ^uint64(0) && offset != res.offset {
 			// Offset control is in use and the offset differs.
 			if offset < res.offset {
-				// The client is confused.
-				// Abort before we break something.
-				return errors.New("unable to rewind ID generation")
+				// This probbably means that 2 clients are running at the same time.
+				// This is fine - just tell the client that is behind what had been dealt with.
+				return ErrIDOffsetDesync{
+					Requested: offset,
+					Base:      res.offset,
+				}
 			}
 
 			// Roll the IDs forward.