try to avoid flooding the logs if a forest goes offline

disable.xqy

@@ -23,6 +23,7 @@ xquery version "1.0-ml";
 import module namespace trb="com.blakeley.task-rebalancer"
   at "lib-trb.xqy" ;
 
+trb:assert-task-server(),
 trb:fatal-set(true())
 
 (: disable :)

enable.xqy

@@ -23,6 +23,7 @@ xquery version "1.0-ml";
 import module namespace trb="com.blakeley.task-rebalancer"
   at "lib-trb.xqy" ;
 
+trb:assert-task-server(),
 trb:fatal-set(false())
 
 (: enable :)

lib-trb.xqy

@@ -38,6 +38,13 @@ declare namespace fs="http://marklogic.com/xdmp/status/forest" ;
 declare variable $FATAL := xdmp:get-server-field(
   'com.blakeley.task-rebalancer.FATAL') ;
 
+(: These error codes indicate a low-level database problem,
+ : which should be treated as fatal.
+ : Used by rebalance.xqy initially.
+ :)
+declare variable $FATAL-CODES := (
+  'XDMP-FORESTNOT') ;
+
 declare variable $HOST := xdmp:host() ;
 
 declare variable $TASKS-COUNT := 0 ;
@@ -47,7 +54,11 @@ declare variable $URI-LAST := () ;
 declare function trb:fatal-set($value as xs:boolean)
  as empty-sequence()
 {
-  xdmp:set-server-field('com.blakeley.task-rebalancer.FATAL', $value)[0],
+  xdmp:log(
+    text { '[trb:fatal-set]',
+      (: Take advantage of set-server-field return value. :)
+      xdmp:set-server-field('com.blakeley.task-rebalancer.FATAL', $value) },
+    'info'),
   xdmp:set($trb:FATAL, $value)
 };
 
@@ -58,6 +69,20 @@ as empty-sequence()
   else error((), 'TRB-FATAL', 'FATAL is set: stopping')
 };
 
+declare function trb:assert-task-server()
+as empty-sequence()
+{
+  (: Assert that this query is running on the task server. :)
+  if (xdmp:server-name(xdmp:server()) eq 'TaskServer') then ()
+  else error(
+    (), 'TRB-SERVERNOT',
+    text {
+      xdmp:server(),
+      xdmp:server-name(xdmp:server()),
+      'This query must run on the Task Server.',
+      'Use xdmp:spawn instead of xdmp:invoke.' })
+};
+
 declare private function trb:uris-start-name(
   $forest as xs:unsignedLong)
 as xs:string
@@ -354,7 +379,7 @@ declare function trb:forests-map(
   $forests as xs:unsignedLong*)
 as map:map
 {
-  xdmp:log(text { 'trb:forest-map', xdmp:describe($forests) }, 'debug'),
+  xdmp:log(text { '[trb:forest-map]', xdmp:describe($forests) }, 'debug'),
   if (empty($forests)) then $m else (
     map:put(
       $m, string($forests[1]),

rebalance.xqy

@@ -30,26 +30,44 @@ declare variable $DOC as node()? := doc($URI) ;
 (: Someday it would be nice to support properties... :)
 declare variable $PROP as node()? := xdmp:document-properties($URI) ;
 
-(: We need a bail-out mechanism to stop the respawns.
- : This document acts as a kill signal.
- : Do not throw an error, since that could flood the server logs.
- :)
-if ($trb:FATAL) then xdmp:log('[rebalance.xqy] FATAL is set: stopping', 'fine')
-else if ($ASSIGNMENT = xdmp:document-forest($URI)) then xdmp:log(
-  text { '[rebalance.xqy]', $URI, 'is already in', $ASSIGNMENT },
-  'fine')
-(: Update the existing document, by overwriting it.
- : explicitly place the document in the correct forest.
+(: In an attempt to catch low-level database errors,
+ : wrap the entire body.
  :)
+try {
+
+  (: We need a bail-out mechanism to stop the respawns.
+   : This document acts as a kill signal.
+   : Do not throw an error, since that could flood the server logs.
+   :)
+  if ($trb:FATAL) then xdmp:log(
+    '[rebalance.xqy] FATAL is set: stopping', 'fine')
+  else if ($ASSIGNMENT = xdmp:document-forest($URI)) then xdmp:log(
+    text { '[rebalance.xqy]', $URI, 'is already in', $ASSIGNMENT },
+    'fine')
+
+  (: Update the existing document, by overwriting it.
+   : explicitly place the document in the correct forest.
+   :)
+  else if (empty($DOC)) then error(
+    (), 'TRB-EMPTY',
+    text { 'No document for', $URI })
+
+  else xdmp:document-insert(
+    $URI,
+    $DOC,
+    xdmp:document-get-permissions($URI),
+    xdmp:document-get-collections($URI),
+    xdmp:document-get-quality($URI),
+    $ASSIGNMENT) }
 
-else if ($DOC) then xdmp:document-insert(
-  $URI,
-  $DOC,
-  xdmp:document-get-permissions($URI),
-  xdmp:document-get-collections($URI),
-  xdmp:document-get-quality($URI),
-  $ASSIGNMENT)
+catch ($ex) {
+  (: Set fatal if the exception looks like a database problem.
+   : This tries to avoid flooding the logs with duplicate error messages.
+   : However these are low-level errors so there are no guarantees.
+   :)
+  if (not($ex/error:code = $trb:FATAL-CODES)) then ()
+  else trb:fatal-set(true()),
+  xdmp:rethrow() }
 
-else error((), 'TRB-UNEXPECTED', ('No document for', $URI))
 
 (: rebalance.xqy :)