From 2f4a2bf98f725d540db9c46747d6234eb872f1b0 Mon Sep 17 00:00:00 2001
From: Nia Weiss <nia.weiss@molecula.com>
Date: Thu, 11 Mar 2021 10:48:12 -0500
Subject: [PATCH] fix ID allocation API after disco

The ID allocation API was broken because the operations were removed from the list allowed in the NORMAL cluster state.
Additionally the operations were set to only run on non-primaries (where they were actually only supposed to run on the primary).
---
 api.go      | 15 +++++++++------
 api_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/api.go b/api.go
index 5815ce21b..d54dd77d6 100644
--- a/api.go
+++ b/api.go
@@ -2174,10 +2174,10 @@ func (api *API) ReserveIDs(key IDAllocKey, session [32]byte, offset uint64, coun
 	snap := topology.NewClusterSnapshot(api.cluster.noder, api.cluster.Hasher, api.cluster.ReplicaN)
 
 	if !snap.IsPrimaryFieldTranslationNode(api.NodeID()) {
-		return api.holder.ida.reserve(key, session, offset, count)
+		return nil, errors.New("cannot reserve IDs on a non-primary node")
 	}
 
-	return nil, errors.New("cannot reserve IDs on a non-primary node")
+	return api.holder.ida.reserve(key, session, offset, count)
 }
 
 func (api *API) CommitIDs(key IDAllocKey, session [32]byte, count uint64) error {
@@ -2189,10 +2189,10 @@ func (api *API) CommitIDs(key IDAllocKey, session [32]byte, count uint64) error
 	snap := topology.NewClusterSnapshot(api.cluster.noder, api.cluster.Hasher, api.cluster.ReplicaN)
 
 	if !snap.IsPrimaryFieldTranslationNode(api.NodeID()) {
-		return api.holder.ida.commit(key, session, count)
+		return errors.New("cannot commit IDs on a non-primary node")
 	}
 
-	return errors.New("cannot commit IDs on a non-primary node")
+	return api.holder.ida.commit(key, session, count)
 }
 
 func (api *API) ResetIDAlloc(index string) error {
@@ -2204,10 +2204,10 @@ func (api *API) ResetIDAlloc(index string) error {
 	snap := topology.NewClusterSnapshot(api.cluster.noder, api.cluster.Hasher, api.cluster.ReplicaN)
 
 	if !snap.IsPrimaryFieldTranslationNode(api.NodeID()) {
-		return api.holder.ida.reset(index)
+		return errors.New("cannot reset IDs on a non-primary node")
 	}
 
-	return errors.New("cannot reset IDs on a non-primary node")
+	return api.holder.ida.reset(index)
 }
 
 // TranslateIndexDB is an internal function to load the index keys database
@@ -2351,4 +2351,7 @@ var methodsNormal = map[apiMethod]struct{}{
 	apiGetTransaction:       {},
 	apiActiveQueries:        {},
 	apiPastQueries:          {},
+	apiIDReserve:            {},
+	apiIDCommit:             {},
+	apiIDReset:              {},
 }
diff --git a/api_test.go b/api_test.go
index c9cb541cf..bbd79f36a 100644
--- a/api_test.go
+++ b/api_test.go
@@ -16,6 +16,7 @@ package pilosa_test
 
 import (
 	"context"
+	"crypto/rand"
 	"fmt"
 	"math"
 	"reflect"
@@ -626,3 +627,45 @@ func TestAPI_ClearFlagForImportAndImportValues(t *testing.T) {
 		panic(fmt.Sprintf("expected %v, observed %v starting acct0 balance", acct0bal, 0))
 	}
 }
+
+func TestAPI_IDAlloc(t *testing.T) {
+	c := test.MustRunCluster(t, 3)
+	defer c.Close()
+
+	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)
+	}
+
+	const toReserve = 2
+
+	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)
+	}
+
+	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)
+	}
+}