Fix bug in inmem ID generation (#312)

Fixes #308

server/datastore/inmem.go

@@ -1,6 +1,7 @@
 package datastore
 
 import (
+	"reflect"
 	"sync"
 
 	"github.com/kolide/kolide-ose/server/kolide"
@@ -8,8 +9,10 @@ import (
 
 type inmem struct {
 	kolide.Datastore
-	Driver               string
-	mtx                  sync.RWMutex
+	Driver  string
+	mtx     sync.RWMutex
+	nextIDs map[interface{}]uint
+
 	users                map[uint]*kolide.User
 	sessions             map[uint]*kolide.Session
 	passwordResets       map[uint]*kolide.PasswordResetRequest
@@ -19,7 +22,8 @@ type inmem struct {
 	queries              map[uint]*kolide.Query
 	packs                map[uint]*kolide.Pack
 	hosts                map[uint]*kolide.Host
-	orginfo              *kolide.OrgInfo
+
+	orginfo *kolide.OrgInfo
 }
 
 func (orm *inmem) Name() string {
@@ -29,6 +33,7 @@ func (orm *inmem) Name() string {
 func (orm *inmem) Migrate() error {
 	orm.mtx.Lock()
 	defer orm.mtx.Unlock()
+	orm.nextIDs = make(map[interface{}]uint)
 	orm.users = make(map[uint]*kolide.User)
 	orm.sessions = make(map[uint]*kolide.Session)
 	orm.passwordResets = make(map[uint]*kolide.PasswordResetRequest)
@@ -64,3 +69,11 @@ func (orm *inmem) getLimitOffsetSliceBounds(opt kolide.ListOptions, length int)
 	}
 	return offset, max
 }
+
+// nextID returns the next ID value that should be used for a struct of the
+// given type
+func (orm *inmem) nextID(val interface{}) uint {
+	valType := reflect.TypeOf(val)
+	orm.nextIDs[valType]++
+	return orm.nextIDs[valType]
+}

server/datastore/inmem_hosts.go

@@ -18,8 +18,7 @@ func (orm *inmem) NewHost(host *kolide.Host) (*kolide.Host, error) {
 		}
 	}
 
-	host.ID = uint(len(orm.hosts) + 1)
-	orm.hosts[host.ID] = host
+	host.ID = orm.nextID(host)
 
 	return host, nil
 }
@@ -122,7 +121,7 @@ func (orm *inmem) EnrollHost(uuid, hostname, ip, platform string, nodeKeySize in
 	}
 
 	if host.ID == 0 {
-		host.ID = uint(len(orm.hosts) + 1)
+		host.ID = orm.nextID(host)
 	}
 	orm.hosts[host.ID] = &host
 

server/datastore/inmem_invites.go

@@ -17,7 +17,7 @@ func (orm *inmem) NewInvite(invite *kolide.Invite) (*kolide.Invite, error) {
 		}
 	}
 
-	invite.ID = uint(len(orm.invites) + 1)
+	invite.ID = orm.nextID(invite)
 	orm.invites[invite.ID] = invite
 	return invite, nil
 }

server/datastore/inmem_labels.go

@@ -20,7 +20,7 @@ func (orm *inmem) NewLabel(label *kolide.Label) (*kolide.Label, error) {
 		}
 	}
 
-	newLabel.ID = uint(len(orm.labels) + 1)
+	newLabel.ID = orm.nextID(label)
 	orm.labels[newLabel.ID] = &newLabel
 
 	return &newLabel, nil
@@ -112,12 +112,12 @@ func (orm *inmem) RecordLabelQueryExecutions(host *kolide.Host, results map[stri
 		if !updated {
 			// Create new execution
 			lqe := kolide.LabelQueryExecution{
-				ID:        uint(len(orm.labelQueryExecutions) + 1),
 				HostID:    host.ID,
 				LabelID:   label.ID,
 				UpdatedAt: t,
 				Matches:   matches,
 			}
+			lqe.ID = orm.nextID(lqe)
 			orm.labelQueryExecutions[lqe.ID] = &lqe
 		}
 	}

server/datastore/inmem_packs.go

@@ -18,7 +18,7 @@ func (orm *inmem) NewPack(pack *kolide.Pack) error {
 		}
 	}
 
-	newPack.ID = uint(len(orm.packs) + 1)
+	newPack.ID = orm.nextID(pack)
 	orm.packs[newPack.ID] = &newPack
 
 	return nil

server/datastore/inmem_password_reset.go

@@ -6,7 +6,7 @@ func (orm *inmem) NewPasswordResetRequest(req *kolide.PasswordResetRequest) (*ko
 	orm.mtx.Lock()
 	defer orm.mtx.Unlock()
 
-	req.ID = uint(len(orm.passwordResets) + 1)
+	req.ID = orm.nextID(req)
 	orm.passwordResets[req.ID] = req
 	return req, nil
 }

server/datastore/inmem_queries.go

@@ -18,7 +18,7 @@ func (orm *inmem) NewQuery(query *kolide.Query) (*kolide.Query, error) {
 		}
 	}
 
-	newQuery.ID = uint(len(orm.queries) + 1)
+	newQuery.ID = orm.nextID(newQuery)
 	orm.queries[newQuery.ID] = &newQuery
 
 	return &newQuery, nil

server/datastore/inmem_sessions.go

@@ -48,7 +48,7 @@ func (orm *inmem) NewSession(session *kolide.Session) (*kolide.Session, error) {
 	orm.mtx.Lock()
 	defer orm.mtx.Unlock()
 
-	session.ID = uint(len(orm.sessions))
+	session.ID = orm.nextID(session)
 	orm.sessions[session.ID] = session
 	if err := orm.MarkSessionAccessed(session); err != nil {
 		return nil, err

server/datastore/inmem_users.go

@@ -16,7 +16,7 @@ func (orm *inmem) NewUser(user *kolide.User) (*kolide.User, error) {
 		}
 	}
 
-	user.ID = uint(len(orm.users) + 1)
+	user.ID = orm.nextID(user)
 	orm.users[user.ID] = user
 
 	return user, nil