dynamic host volumes: fix unblocking bug in state store

The `HostVolumeByID` state store method didn't add a watch channel to the
watchset, which meant that it would never unblock. The tests missed this because
they were racy, so move the updates for unblocking tests into a `time.After`
call to ensure the queries are blocked before the update happens.

nomad/host_volume_endpoint_test.go

@@ -177,9 +177,14 @@ func TestHostVolumeEndpoint_CreateRegisterGetDelete(t *testing.T) {
 			volCh <- getResp.Volume
 		}()
 
-		var registerResp structs.HostVolumeRegisterResponse
-		err = msgpackrpc.CallWithCodec(codec, "HostVolume.Register", registerReq, &registerResp)
-		must.NoError(t, err)
+		// re-register the volume long enough later that we can be sure we won't
+		// win a race with the get RPC goroutine
+		time.AfterFunc(200*time.Millisecond, func() {
+			codec := rpcClient(t, srv)
+			var registerResp structs.HostVolumeRegisterResponse
+			err = msgpackrpc.CallWithCodec(codec, "HostVolume.Register", registerReq, &registerResp)
+			must.NoError(t, err)
+		})
 
 		select {
 		case <-ctx.Done():
@@ -480,8 +485,14 @@ func TestHostVolumeEndpoint_List(t *testing.T) {
 			respCh <- &listResp
 		}()
 
-		err = msgpackrpc.CallWithCodec(codec, "HostVolume.Register", registerReq, &registerResp)
-		must.NoError(t, err)
+		// re-register the volume long enough later that we can be sure we won't
+		// win a race with the get RPC goroutine
+		time.AfterFunc(200*time.Millisecond, func() {
+			codec := rpcClient(t, srv)
+			var registerResp structs.HostVolumeRegisterResponse
+			err = msgpackrpc.CallWithCodec(codec, "HostVolume.Register", registerReq, &registerResp)
+			must.NoError(t, err)
+		})
 
 		select {
 		case <-ctx.Done():

nomad/state/state_store_host_volumes.go

@@ -13,10 +13,12 @@ import (
 // HostVolumeByID retrieve a specific host volume
 func (s *StateStore) HostVolumeByID(ws memdb.WatchSet, ns, id string, withAllocs bool) (*structs.HostVolume, error) {
 	txn := s.db.ReadTxn()
-	obj, err := txn.First(TableHostVolumes, indexID, ns, id)
+	watchCh, obj, err := txn.FirstWatch(TableHostVolumes, indexID, ns, id)
 	if err != nil {
 		return nil, err
 	}
+	ws.Add(watchCh)
+
 	if obj == nil {
 		return nil, nil
 	}