[NFC][Stack Traces] Do not build a functor where we can use the translation map directly

PiperOrigin-RevId: 356403200
Change-Id: I6ae0231a08bcc24b103f3fa8f95233aa2116a952
This commit is contained in:
George Karpenkov 2021-02-08 18:45:17 -08:00 committed by TensorFlower Gardener
parent 2c596172d2
commit 969371e28f
5 changed files with 20 additions and 31 deletions

View File

@ -482,6 +482,7 @@ cc_library(
],
deps = [
"//tensorflow/core/platform:stack_frame",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],

View File

@ -19,20 +19,22 @@ limitations under the License.
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/strings/match.h"
#include "absl/types/optional.h"
#include "tensorflow/core/platform/stack_frame.h"
namespace tensorflow {
using SourceLoc = std::tuple<std::string, int>;
// Maps filename/line_no combination into a stack frame.
using StackTraceMap =
std::function<absl::optional<StackFrame>(std::pair<const char*, int>)>;
using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
// Returns "true" on filenames which should be skipped.
using StackTraceFilter = std::function<bool(const char*)>;
using ToStackFramesFunctor = std::vector<StackFrame>(int, const StackTraceMap&,
using ToStackFramesFunctor = std::vector<StackFrame>(int, const SourceMap&,
const StackTraceFilter&,
bool, int);
@ -54,11 +56,12 @@ class ManagedStackTrace {
: id_(id), to_stack_frames_(to_stack_frames) {}
// Returns stack trace as a vector of `StackFrame`s.
std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
const StackTraceFilter& filtered = {},
bool reverse_traversal = false,
int limit = -1) const {
return to_stack_frames_(id_, mapper, filtered, reverse_traversal, limit);
return to_stack_frames_(id_, source_map, filtered, reverse_traversal,
limit);
}
private:

View File

@ -42,7 +42,7 @@ const char* GetPythonString(PyObject* o) {
namespace tensorflow {
std::vector<StackFrame> StackTrace::ToStackFrames(
const StackTraceMap& mapper, const StackTraceFilter& filtered,
const SourceMap& source_map, const StackTraceFilter& filtered,
bool reverse_traversal, int limit) const {
DCheckPyGilStateForStackTrace();
std::vector<StackFrame> result;
@ -61,11 +61,9 @@ std::vector<StackFrame> StackTrace::ToStackFrames(
continue;
}
absl::optional<StackFrame> mapped =
mapper ? mapper(std::make_pair(file_name, line_number)) : absl::nullopt;
if (mapped) {
result.push_back(*mapped);
const auto it = source_map.find(std::make_tuple(file_name, line_number));
if (it != source_map.end()) {
result.push_back(it->second);
} else {
result.emplace_back(StackFrame{file_name, line_number,
GetPythonString(code_obj.first->co_name)});

View File

@ -88,12 +88,12 @@ class StackTrace final {
}
// Returns a structured representation of the captured stack trace.
// `mapper` provides a custom mapping for translating stack frames, `filter`
// returns `true` for the stack frames which should be omitted.
// `source_map` provides a custom mapping for translating stack frames,
// `filter` returns `true` for the stack frames which should be omitted.
//
// `reverse_traversal` changes the traversal order of the stack trace, and
// `limit` bounds the number of returned frames (after filtering).
std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
const StackTraceFilter& filtered = {},
bool reverse_traversal = false,
int limit = -1) const;
@ -149,11 +149,11 @@ extern StackTraceManager* const stack_trace_manager;
// Converts the ManagedStackTrace (identified by ID) to a vector of stack
// frames.
inline std::vector<StackFrame> ManagedStackTraceToStackFrames(
int id, const StackTraceMap& mapper, const StackTraceFilter& filtered,
int id, const SourceMap& source_map, const StackTraceFilter& filtered,
bool reverse_traversal, int limit) {
PyGILState_STATE gstate = PyGILState_Ensure();
std::vector<StackFrame> result = stack_trace_manager->Get(id)->ToStackFrames(
mapper, filtered, reverse_traversal, limit);
source_map, filtered, reverse_traversal, limit);
PyGILState_Release(gstate);
return result;
}

View File

@ -54,10 +54,6 @@ namespace {
namespace py = pybind11;
using SourceLoc = std::tuple<std::string, int>;
using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
using StringSet = absl::flat_hash_set<std::string>;
// Python wrapper for a SourceMap.
@ -149,8 +145,7 @@ class StackTraceWrapper : public AbstractStackTrace {
PyGILState_STATE state = PyGILState_Ensure();
stack_frames_cache_ = captured_.ToStackFrames(
[&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
[&](const char* f) { return StackTraceFiltering(f); });
*source_map_, [&](const char* f) { return StackTraceFiltering(f); });
stack_frames_cache_->pop_back(); // Drop last stack frame.
PyGILState_Release(state);
return *stack_frames_cache_;
@ -163,7 +158,7 @@ class StackTraceWrapper : public AbstractStackTrace {
PyGILState_STATE state = PyGILState_Ensure();
std::vector<StackFrame> last_frame = captured_.ToStackFrames(
[&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
*source_map_,
[&](const char* file_name) {
return StackTraceFiltering(file_name) ||
IsInternalFrameForFilename(file_name);
@ -226,14 +221,6 @@ class StackTraceWrapper : public AbstractStackTrace {
});
}
absl::optional<StackFrame> StackTraceMapping(SourceLoc loc) const {
if (source_map_->contains(loc)) {
return source_map_->at(loc);
}
return absl::nullopt;
}
bool StackTraceFiltering(const char* file_name) const {
return filter_->contains(file_name);
}