diff --git a/src/Brofiler.cc b/src/Brofiler.cc index 5080e25f54..5603040786 100644 --- a/src/Brofiler.cc +++ b/src/Brofiler.cc @@ -15,6 +15,17 @@ Brofiler::Brofiler() Brofiler::~Brofiler() { + for ( auto& s : stmts ) + Unref(s); + } + +void Brofiler::AddStmt(Stmt* s) + { + if ( ignoring != 0 ) + return; + + ::Ref(s); + stmts.push_back(s); } bool Brofiler::ReadStats() @@ -109,7 +120,7 @@ bool Brofiler::WriteStats() return false; } - for ( list::const_iterator it = stmts.begin(); + for ( list::const_iterator it = stmts.begin(); it != stmts.end(); ++it ) { ODesc location_info; diff --git a/src/Brofiler.h b/src/Brofiler.h index b9a0a00244..504b6e324f 100644 --- a/src/Brofiler.h +++ b/src/Brofiler.h @@ -38,13 +38,13 @@ public: void IncIgnoreDepth() { ignoring++; } void DecIgnoreDepth() { ignoring--; } - void AddStmt(const Stmt* s) { if ( ignoring == 0 ) stmts.push_back(s); } + void AddStmt(Stmt* s); private: /** * The current, global Brofiler instance creates this list at parse-time. */ - list stmts; + list stmts; /** * Indicates whether new statments will not be considered as part of diff --git a/src/Expr.cc b/src/Expr.cc index 9eba79f02c..f5db8ff0fb 100644 --- a/src/Expr.cc +++ b/src/Expr.cc @@ -4339,6 +4339,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr arg_ing, // Install a dummy version of the function globally for use only // when broker provides a closure. + ::Ref(ingredients->body); BroFunc* dummy_func = new BroFunc( ingredients->id, ingredients->body, @@ -4378,6 +4379,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr arg_ing, dummy_func->SetName(my_name.c_str()); Val* v = new Val(dummy_func); + Unref(dummy_func); id->SetVal(v); // id will unref v when its done. id->SetType(ingredients->id->Type()->Ref()); id->SetConst(); @@ -4385,6 +4387,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr arg_ing, Val* LambdaExpr::Eval(Frame* f) const { + ::Ref(ingredients->body); BroFunc* lamb = new BroFunc( ingredients->id, ingredients->body, @@ -4398,7 +4401,9 @@ Val* LambdaExpr::Eval(Frame* f) const // Allows for lookups by the receiver. lamb->SetName(my_name.c_str()); - return new Val(lamb); + auto rval = new Val(lamb); + Unref(lamb); + return rval; } void LambdaExpr::ExprDescribe(ODesc* d) const diff --git a/src/Frame.cc b/src/Frame.cc index f4eb49cbac..249d988b7b 100644 --- a/src/Frame.cc +++ b/src/Frame.cc @@ -31,20 +31,54 @@ Frame::Frame(int arg_size, const BroFunc* func, const val_list* fn_args) Frame::~Frame() { + for ( auto& func : functions_with_closure_frame_reference ) + { + func->StrengthenClosureReference(this); + Unref(func); + } + // Deleting a Frame that is a view is a no-op. Unref(trigger); - Unref(closure); + + if ( ! weak_closure_ref ) + Unref(closure); for ( auto& i : outer_ids ) Unref(i); Release(); + + delete [] weak_refs; } -void Frame::SetElement(int n, Val* v) +void Frame::AddFunctionWithClosureRef(BroFunc* func) { - Unref(frame[n]); + ::Ref(func); + functions_with_closure_frame_reference.emplace_back(func); + } + +void Frame::SetElement(int n, Val* v, bool weak_ref) + { + UnrefElement(n); frame[n] = v; + + if ( weak_ref ) + { + if ( ! weak_refs ) + { + weak_refs = new bool[size]; + + for ( auto i = 0; i < size; ++i ) + weak_refs[i] = false; + } + + weak_refs[n] = true; + } + else + { + if ( weak_refs ) + weak_refs[n] = false; + } } void Frame::SetElement(const ID* id, Val* v) @@ -62,8 +96,15 @@ void Frame::SetElement(const ID* id, Val* v) if ( offset_map.size() ) { auto where = offset_map.find(std::string(id->Name())); + if ( where != offset_map.end() ) - SetElement(where->second, v); + { + // Need to add a Ref to 'v' since the SetElement() for + // id->Offset() below is otherwise responsible for keeping track + // of the implied reference count of the passed-in 'v' argument. + // i.e. if we end up storing it twice, we need an addition Ref. + SetElement(where->second, v->Ref()); + } } SetElement(id->Offset(), v); @@ -92,7 +133,7 @@ void Frame::Reset(int startIdx) { for ( int i = startIdx; i < size; ++i ) { - Unref(frame[i]); + UnrefElement(i); frame[i] = 0; } } @@ -100,7 +141,7 @@ void Frame::Reset(int startIdx) void Frame::Release() { for ( int i = 0; i < size; ++i ) - Unref(frame[i]); + UnrefElement(i); delete [] frame; } @@ -145,7 +186,34 @@ Frame* Frame::Clone() const return other; } -Frame* Frame::SelectiveClone(const id_list& selection) const +static bool val_is_func(Val* v, BroFunc* func) + { + if ( v->Type()->Tag() != TYPE_FUNC ) + return false; + + return v->AsFunc() == func; + } + +static Val* clone_if_not_func(Val** frame, int offset, BroFunc* func, + Frame* other) + { + auto v = frame[offset]; + + if ( ! v ) + return nullptr; + + if ( val_is_func(v, func) ) + { + other->SetElement(offset, v, true); + return v; + } + + auto rval = v->Clone(); + other->SetElement(offset, rval); + return rval; + } + +Frame* Frame::SelectiveClone(const id_list& selection, BroFunc* func) const { if ( selection.length() == 0 ) return nullptr; @@ -171,7 +239,7 @@ Frame* Frame::SelectiveClone(const id_list& selection) const auto where = offset_map.find(std::string(id->Name())); if ( where != offset_map.end() ) { - other->frame[where->second] = frame[where->second]->Clone(); + clone_if_not_func(frame, where->second, func, other); continue; } } @@ -179,7 +247,7 @@ Frame* Frame::SelectiveClone(const id_list& selection) const if ( ! frame[id->Offset()] ) reporter->InternalError("Attempted to clone an id ('%s') with no associated value.", id->Name()); - other->frame[id->Offset()] = frame[id->Offset()]->Clone(); + clone_if_not_func(frame, id->Offset(), func, other); } /** @@ -379,6 +447,7 @@ std::pair Frame::Unserialize(const broker::vector& data) // Frame takes ownership of unref'ing elements in outer_ids rf->outer_ids = std::move(outer_ids); rf->closure = closure; + rf->weak_closure_ref = false; for ( int i = 0; i < frame_size; ++i ) { @@ -437,7 +506,7 @@ void Frame::CaptureClosure(Frame* c, id_list arg_outer_ids) closure = c; if ( closure ) - Ref(closure); + weak_closure_ref = true; /** * Want to capture closures by copy? diff --git a/src/Frame.h b/src/Frame.h index efb1e2c398..133ab45978 100644 --- a/src/Frame.h +++ b/src/Frame.h @@ -44,8 +44,11 @@ public: * * @param n the index to set * @param v the value to set it to + * @param weak_ref whether the frame owns the value and should unref + * it upon destruction. Used to break circular references between + * lambda functions and closure frames. */ - void SetElement(int n, Val* v); + void SetElement(int n, Val* v, bool weak_ref = false); /** * Associates *id* and *v* in the frame. Future lookups of @@ -149,7 +152,7 @@ public: * *selection* have been cloned. All other values are made to be * null. */ - Frame* SelectiveClone(const id_list& selection) const; + Frame* SelectiveClone(const id_list& selection, BroFunc* func) const; /** * Serializes the Frame into a Broker representation. @@ -215,8 +218,28 @@ public: void SetDelayed() { delayed = true; } bool HasDelayed() const { return delayed; } + /** + * Track a new function that refers to this frame for use as a closure. + * This frame's destructor will then upgrade that functions reference + * from weak to strong (by making a copy). The initial use of + * weak references prevents unbreakable circular references that + * otherwise cause memory leaks. + */ + void AddFunctionWithClosureRef(BroFunc* func); private: + + /** + * Unrefs the value at offset 'n' frame unless it's a weak reference. + */ + void UnrefElement(int n) + { + if ( weak_refs && weak_refs[n] ) + return; + + Unref(frame[n]); + } + /** Have we captured this id? */ bool IsOuterID(const ID* in) const; @@ -242,8 +265,13 @@ private: /** Associates ID's offsets with values. */ Val** frame; + /** Values that are weakly referenced by the frame. Used to + * prevent circular reference memory leaks in lambda/closures */ + bool* weak_refs = nullptr; + /** The enclosing frame of this frame. */ Frame* closure; + bool weak_closure_ref = false; /** ID's used in this frame from the enclosing frame. */ id_list outer_ids; @@ -268,6 +296,8 @@ private: Trigger* trigger; const CallExpr* call; bool delayed; + + std::vector functions_with_closure_frame_reference; }; /** diff --git a/src/Func.cc b/src/Func.cc index 747029e4f8..fb90e2e3d2 100644 --- a/src/Func.cc +++ b/src/Func.cc @@ -132,6 +132,7 @@ Func* Func::DoClone() { // By default, ok just to return a reference. Func does not have any state // that is different across instances. + ::Ref(this); return this; } @@ -286,7 +287,9 @@ BroFunc::~BroFunc() { std::for_each(bodies.begin(), bodies.end(), [](Body& b) { Unref(b.stmts); }); - Unref(closure); + + if ( ! weak_closure_ref ) + Unref(closure); } int BroFunc::IsPure() const @@ -490,14 +493,35 @@ void BroFunc::AddClosure(id_list ids, Frame* f) SetClosureFrame(f); } +bool BroFunc::StrengthenClosureReference(Frame* f) + { + if ( closure != f ) + return false; + + if ( ! weak_closure_ref ) + return false; + + closure = closure->SelectiveClone(outer_ids, this); + weak_closure_ref = false; + return true; + } + void BroFunc::SetClosureFrame(Frame* f) { if ( closure ) reporter->InternalError("Tried to override closure for BroFunc %s.", Name()); + // Have to use weak references initially because otherwise Ref'ing the + // original frame creates a circular reference: the function holds a + // reference to the frame and the frame contains a reference to this + // function value. And we can't just do a shallow clone of the frame + // up front because the closure semantics in Zeek allow mutating + // the outer frame. + closure = f; - Ref(closure); + weak_closure_ref = true; + f->AddFunctionWithClosureRef(this); } bool BroFunc::UpdateClosure(const broker::vector& data) @@ -510,9 +534,10 @@ bool BroFunc::UpdateClosure(const broker::vector& data) if ( new_closure ) new_closure->SetFunction(this); - if ( closure ) + if ( ! weak_closure_ref ) Unref(closure); + weak_closure_ref = false; closure = new_closure; return true; @@ -528,7 +553,8 @@ Func* BroFunc::DoClone() CopyStateInto(other); other->frame_size = frame_size; - other->closure = closure ? closure->SelectiveClone(outer_ids) : nullptr; + other->closure = closure ? closure->SelectiveClone(outer_ids, this) : nullptr; + other->weak_closure_ref = false; other->outer_ids = outer_ids; return other; @@ -811,3 +837,68 @@ bool check_built_in_call(BuiltinFunc* f, CallExpr* call) return true; } + +// Gets a function's priority from its Scope's attributes. Errors if it sees any +// problems. +static int get_func_priority(const attr_list& attrs) + { + int priority = 0; + + for ( const auto& a : attrs ) + { + if ( a->Tag() == ATTR_DEPRECATED ) + continue; + + if ( a->Tag() != ATTR_PRIORITY ) + { + a->Error("illegal attribute for function body"); + continue; + } + + Val* v = a->AttrExpr()->Eval(0); + if ( ! v ) + { + a->Error("cannot evaluate attribute expression"); + continue; + } + + if ( ! IsIntegral(v->Type()->Tag()) ) + { + a->Error("expression is not of integral type"); + continue; + } + + priority = v->InternalInt(); + } + + return priority; + } + +function_ingredients::function_ingredients(Scope* scope, Stmt* body) + { + frame_size = scope->Length(); + inits = scope->GetInits(); + + this->scope = scope; + ::Ref(this->scope); + id = scope->ScopeID(); + ::Ref(id); + + auto attrs = scope->Attrs(); + + priority = (attrs ? get_func_priority(*attrs) : 0); + this->body = body; + ::Ref(this->body); + } + +function_ingredients::~function_ingredients() + { + Unref(id); + Unref(body); + Unref(scope); + + for ( const auto& i : *inits ) + Unref(i); + + delete inits; + } diff --git a/src/Func.h b/src/Func.h index fc565d61ab..802c4509da 100644 --- a/src/Func.h +++ b/src/Func.h @@ -114,6 +114,12 @@ public: */ bool UpdateClosure(const broker::vector& data); + /** + * If the function's closure is a weak reference to the given frame, + * upgrade to a strong reference of a shallow clone of that frame. + */ + bool StrengthenClosureReference(Frame* f); + /** * Serializes this function's closure. * @@ -154,6 +160,7 @@ private: id_list outer_ids; // The frame the BroFunc was initialized in. Frame* closure = nullptr; + bool weak_closure_ref = false; }; typedef Val* (*built_in_func)(Frame* frame, val_list* args); @@ -192,13 +199,20 @@ struct CallInfo { // Struct that collects all the specifics defining a Func. Used for BroFuncs // with closures. struct function_ingredients { + + // Gathers all of the information from a scope and a function body needed + // to build a function. + function_ingredients(Scope* scope, Stmt* body); + + ~function_ingredients(); + ID* id; Stmt* body; id_list* inits; int frame_size; int priority; Scope* scope; - }; +}; extern vector call_stack; diff --git a/src/Val.cc b/src/Val.cc index 53aecb0bbb..40533ba748 100644 --- a/src/Val.cc +++ b/src/Val.cc @@ -120,7 +120,12 @@ Val* Val::DoClone(CloneState* state) // Derived classes are responsible for this. Exception: // Functions and files. There aren't any derived classes. if ( type->Tag() == TYPE_FUNC ) - return new Val(AsFunc()->DoClone()); + { + auto c = AsFunc()->DoClone(); + auto rval = new Val(c); + Unref(c); + return rval; + } if ( type->Tag() == TYPE_FILE ) { diff --git a/src/Var.cc b/src/Var.cc index f7e15041e8..212ee58a54 100644 --- a/src/Var.cc +++ b/src/Var.cc @@ -416,13 +416,30 @@ public: : scope(s) { } virtual TraversalCode PreExpr(const Expr*); + virtual TraversalCode PostExpr(const Expr*); Scope* scope; vector outer_id_references; + int lambda_depth = 0; + // Note: think we really ought to toggle this to false to prevent + // considering locals within inner-lambdas as "outer", but other logic + // for "selective cloning" and locating IDs in the closure chain may + // depend on current behavior and also needs to be changed. + bool search_inner_lambdas = true; }; TraversalCode OuterIDBindingFinder::PreExpr(const Expr* expr) { + if ( expr->Tag() == EXPR_LAMBDA ) + ++lambda_depth; + + if ( lambda_depth > 0 && ! search_inner_lambdas ) + // Don't inspect the bodies of inner lambdas as they will have their + // own traversal to find outer IDs and we don't want to detect + // references to local IDs inside and accidentally treat them as + // "outer" since they can't be found in current scope. + return TC_CONTINUE; + if ( expr->Tag() != EXPR_NAME ) return TC_CONTINUE; @@ -438,45 +455,20 @@ TraversalCode OuterIDBindingFinder::PreExpr(const Expr* expr) return TC_CONTINUE; } -// Gets a function's priority from its Scope's attributes. Errors if it sees any -// problems. -static int get_func_priotity(const attr_list& attrs) +TraversalCode OuterIDBindingFinder::PostExpr(const Expr* expr) { - int priority = 0; - - for ( const auto& a : attrs ) + if ( expr->Tag() == EXPR_LAMBDA ) { - if ( a->Tag() == ATTR_DEPRECATED ) - continue; - - if ( a->Tag() != ATTR_PRIORITY ) - { - a->Error("illegal attribute for function body"); - continue; - } - - Val* v = a->AttrExpr()->Eval(0); - if ( ! v ) - { - a->Error("cannot evaluate attribute expression"); - continue; - } - - if ( ! IsIntegral(v->Type()->Tag()) ) - { - a->Error("expression is not of integral type"); - continue; - } - - priority = v->InternalInt(); + --lambda_depth; + assert(lambda_depth >= 0); } - return priority; + return TC_CONTINUE; } void end_func(Stmt* body) { - std::unique_ptr ingredients = gather_function_ingredients(pop_scope(), body); + auto ingredients = std::make_unique(pop_scope(), body); if ( streq(ingredients->id->Name(), "anonymous-function") ) { @@ -508,24 +500,10 @@ void end_func(Stmt* body) } ingredients->id->ID_Val()->AsFunc()->SetScope(ingredients->scope); - } - -std::unique_ptr gather_function_ingredients(Scope* scope, Stmt* body) - { - auto ingredients = std::make_unique(); - - ingredients->frame_size = scope->Length(); - ingredients->inits = scope->GetInits(); - - ingredients->scope = scope; - ingredients->id = scope->ScopeID(); - - auto attrs = scope->Attrs(); - - ingredients->priority = (attrs ? get_func_priotity(*attrs) : 0); - ingredients->body = body; - - return ingredients; + // Note: ideally, something would take ownership of this memory until the + // end of script execution, but that's essentially the same as the + // lifetime of the process at the moment, so ok to "leak" it. + ingredients.release(); } Val* internal_val(const char* name) @@ -548,7 +526,14 @@ id_list gather_outer_ids(Scope* scope, Stmt* body) id_list idl ( cb.outer_id_references.size() ); for ( size_t i = 0; i < cb.outer_id_references.size(); ++i ) - idl.append(cb.outer_id_references[i]->Id()); + { + auto id = cb.outer_id_references[i]->Id(); + + if ( idl.is_member(id) ) + continue; + + idl.append(id); + } return idl; } diff --git a/src/Var.h b/src/Var.h index f90250a3b1..5ae6720b62 100644 --- a/src/Var.h +++ b/src/Var.h @@ -26,11 +26,6 @@ extern void begin_func(ID* id, const char* module_name, function_flavor flavor, int is_redef, FuncType* t, attr_list* attrs = nullptr); extern void end_func(Stmt* body); -// Gathers all of the information from a scope and a function body needed to -// build a function and collects it into a function_ingredients struct. -// Gathered elements are not refeed. -extern std::unique_ptr gather_function_ingredients(Scope* scope, Stmt* body); - // Gather all IDs referenced inside a body that aren't part of a given scope. extern id_list gather_outer_ids(Scope* scope, Stmt* body); diff --git a/src/parse.y b/src/parse.y index 6e8b4dc327..c409942cca 100644 --- a/src/parse.y +++ b/src/parse.y @@ -1234,7 +1234,7 @@ anonymous_function: // a lambda expression. // Gather the ingredients for a BroFunc from the current scope - std::unique_ptr ingredients = gather_function_ingredients(current_scope(), $5); + auto ingredients = std::make_unique(current_scope(), $5); id_list outer_ids = gather_outer_ids(pop_scope(), $5); $$ = new LambdaExpr(std::move(ingredients), std::move(outer_ids));