mol-model/query timeout

src/mol-model/structure/query/context.ts

@@ -5,6 +5,7 @@
  */
 
 import { Structure, StructureElement } from '../structure';
+import { now } from 'mol-task';
 
 export interface QueryContextView {
     readonly element: StructureElement;
@@ -14,6 +15,8 @@ export interface QueryContextView {
 export class QueryContext implements QueryContextView {
     private currentElementStack: StructureElement[] = [];
     private currentStructureStack: Structure[] = [];
+    private timeCreated = now();
+    private timeoutMs: number;
 
     readonly inputStructure: Structure;
 
@@ -40,8 +43,18 @@ export class QueryContext implements QueryContextView {
         else (this.currentStructure as Structure) = void 0 as any;
     }
 
-    constructor(structure: Structure) {
+    throwIfTimedOut() {
+        if (this.timeoutMs === 0) return;
+        if (now() - this.timeCreated > this.timeoutMs) {
+            throw new Error(`The query took too long to execute (> ${this.timeoutMs / 1000}s).`);
+        }
+    }
+
+    // todo timeout
+
+    constructor(structure: Structure, timeoutMs = 0) {
         this.inputStructure = structure;
+        this.timeoutMs = timeoutMs;
     }
 }
 

src/mol-model/structure/query/queries/combinators.ts

@@ -11,8 +11,9 @@ export function merge(queries: ArrayLike<StructureQuery>): StructureQuery {
     return ctx => {
         const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
         for (let i = 0; i < queries.length; i++) {
-            StructureSelection.forEach(queries[i](ctx), s => {
+            StructureSelection.forEach(queries[i](ctx), (s, j) => {
                 ret.add(s);
+                if (i % 100) ctx.throwIfTimedOut();
             });
         }
         return ret.getSelection();

src/mol-model/structure/query/queries/filters.ts

@@ -16,9 +16,10 @@ export function pick(query: StructureQuery, pred: QueryPredicate): StructureQuer
         const sel = query(ctx);
         const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
         ctx.pushCurrentElement();
-        StructureSelection.forEach(sel, s => {
+        StructureSelection.forEach(sel, (s, i) => {
             ctx.currentStructure = s;
             if (pred(ctx)) ret.add(s);
+            if (i % 100) ctx.throwIfTimedOut();
         });
         ctx.popCurrentStructure();
         return ret.getSelection();
@@ -44,6 +45,8 @@ export function getCurrentStructureProperties(ctx: QueryContext, props: UnitType
             l.element = elements[j];
             set.add(fn(ctx));
         }
+
+        ctx.throwIfTimedOut();
     }
     ctx.popCurrentElement();
     return set;
@@ -54,9 +57,11 @@ function getSelectionProperties(ctx: QueryContext, query: StructureQuery, props:
 
     const sel = query(ctx);
     ctx.pushCurrentElement();
-    StructureSelection.forEach(sel, s => {
+    StructureSelection.forEach(sel, (s, i) => {
         ctx.currentStructure = s;
         getCurrentStructureProperties(ctx, props, set);
+
+        if (i % 10) ctx.throwIfTimedOut();
     });
     ctx.popCurrentElement();
     return set;
@@ -69,12 +74,14 @@ export function withSameAtomProperties(query: StructureQuery, propertySource: St
 
         const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
         ctx.pushCurrentStructure();
-        StructureSelection.forEach(sel, s => {
+        StructureSelection.forEach(sel, (s, i) => {
             ctx.currentStructure = s;
             const currentProps = getCurrentStructureProperties(ctx, props, new Set());
             if (isSuperset(currentProps, propSet)) {
                 ret.add(s);
             }
+
+            if (i % 10) ctx.throwIfTimedOut();
         });
         ctx.popCurrentStructure();
         return ret.getSelection();
@@ -86,8 +93,9 @@ export function areIntersectedBy(query: StructureQuery, by: StructureQuery): Str
         const mask = StructureSelection.unionStructure(by(ctx));
         const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
 
-        StructureSelection.forEach(query(ctx), s => {
+        StructureSelection.forEach(query(ctx), (s, i) => {
             if (structureAreIntersecting(mask, s)) ret.add(s);
+            if (i % 10) ctx.throwIfTimedOut();
         });
         return ret.getSelection();
     };

src/mol-model/structure/query/queries/generators.ts

@@ -60,6 +60,8 @@ function atomGroupsLinear(atomTest: QueryPredicate): StructureQuery {
                 if (atomTest(ctx)) builder.addElement(l.element);
             }
             builder.commitUnit();
+
+            ctx.throwIfTimedOut();
         }
         ctx.popCurrentElement();
         return StructureSelection.Singletons(inputStructure, builder.getStructure());
@@ -105,6 +107,8 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
                 }
             }
             builder.commitUnit();
+
+            ctx.throwIfTimedOut();
         }
         ctx.popCurrentElement();
         return StructureSelection.Singletons(inputStructure, builder.getStructure());
@@ -146,6 +150,8 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
                     }
                 }
             }
+
+            ctx.throwIfTimedOut();
         }
         ctx.popCurrentElement();
         return builder.getSelection();

src/mol-model/structure/query/queries/modifiers.ts

@@ -10,8 +10,9 @@ import { StructureQuery } from '../query';
 import { StructureSelection } from '../selection';
 import { UniqueStructuresBuilder } from '../utils/builders';
 import { StructureUniqueSubsetBuilder } from '../../structure/util/unique-subset-builder';
+import { QueryContext } from '../context';
 
-function getWholeResidues(source: Structure, structure: Structure) {
+function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) {
     const builder = source.subsetBuilder(true);
     for (const unit of structure.units) {
         if (unit.kind !== Unit.Kind.Atomic) {
@@ -32,6 +33,8 @@ function getWholeResidues(source: Structure, structure: Structure) {
             }
         }
         builder.commitUnit();
+
+        ctx.throwIfTimedOut();
     }
     return builder.getStructure();
 }
@@ -40,11 +43,11 @@ export function wholeResidues(query: StructureQuery, isFlat: boolean): Structure
     return ctx => {
         const inner = query(ctx);
         if (StructureSelection.isSingleton(inner)) {
-            return StructureSelection.Singletons(ctx.inputStructure, getWholeResidues(ctx.inputStructure, inner.structure));
+            return StructureSelection.Singletons(ctx.inputStructure, getWholeResidues(ctx, ctx.inputStructure, inner.structure));
         } else {
             const builder = new UniqueStructuresBuilder(ctx.inputStructure);
             for (const s of inner.structures) {
-                builder.add(getWholeResidues(ctx.inputStructure, s));
+                builder.add(getWholeResidues(ctx, ctx.inputStructure, s));
             }
             return builder.getSelection();
         }
@@ -60,7 +63,7 @@ export interface IncludeSurroundingsParams {
     wholeResidues?: boolean
 }
 
-function getIncludeSurroundings(source: Structure, structure: Structure, params: IncludeSurroundingsParams) {
+function getIncludeSurroundings(ctx: QueryContext, source: Structure, structure: Structure, params: IncludeSurroundingsParams) {
     const builder = new StructureUniqueSubsetBuilder(source);
     const lookup = source.lookup3d;
     const r = params.radius;
@@ -72,21 +75,23 @@ function getIncludeSurroundings(source: Structure, structure: Structure, params:
             const e = elements[i];
             lookup.findIntoBuilder(x(e), y(e), z(e), r, builder);
         }
+
+        ctx.throwIfTimedOut();
     }
-    return !!params.wholeResidues ? getWholeResidues(source, builder.getStructure()) : builder.getStructure();
+    return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure();
 }
 
 export function includeSurroundings(query: StructureQuery, params: IncludeSurroundingsParams): StructureQuery {
     return ctx => {
         const inner = query(ctx);
         if (StructureSelection.isSingleton(inner)) {
-            const surr = getIncludeSurroundings(ctx.inputStructure, inner.structure, params);
+            const surr = getIncludeSurroundings(ctx, ctx.inputStructure, inner.structure, params);
             const ret = StructureSelection.Singletons(ctx.inputStructure, surr);
             return ret;
         } else {
             const builder = new UniqueStructuresBuilder(ctx.inputStructure);
             for (const s of inner.structures) {
-                builder.add(getIncludeSurroundings(ctx.inputStructure, s, params));
+                builder.add(getIncludeSurroundings(ctx, ctx.inputStructure, s, params));
             }
             return builder.getSelection();
         }

src/mol-model/structure/query/query.ts

@@ -10,8 +10,8 @@ import { QueryContext } from './context';
 
 interface StructureQuery { (ctx: QueryContext): StructureSelection }
 namespace StructureQuery {
-    export function run(query: StructureQuery, structure: Structure) {
-        return query(new QueryContext(structure))
+    export function run(query: StructureQuery, structure: Structure, timeoutMs = 0) {
+        return query(new QueryContext(structure, timeoutMs));
     }
 }
 

src/servers/model/server/query.ts

@@ -32,11 +32,12 @@ export async function resolveJob(job: Job, writer: Writer) {
     const wrappedStructure = await getStructure(job);
 
     perf.start('query');
+    // TODO: encode errors that happen past this point as CIF rather than just 404
     const structure = job.queryDefinition.structureTransform
         ? await job.queryDefinition.structureTransform(job.normalizedParams, wrappedStructure.structure)
         : wrappedStructure.structure;
     const query = job.queryDefinition.query(job.normalizedParams, structure);
-    const result = StructureSelection.unionStructure(StructureQuery.run(query, structure));
+    const result = StructureSelection.unionStructure(StructureQuery.run(query, structure, Config.maxQueryTimeInMs));
     perf.end('query');
 
     ConsoleLogger.logId(job.id, 'Query', 'Query finished.');