Skip to content

Commit e97e80c

Browse files
gf712gf712
committed
Address review comment
1 parent 81b0d4f commit e97e80c

3 files changed

Lines changed: 106 additions & 62 deletions

File tree

‎integration/tests/buffered_writer.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,4 +137,24 @@ def write(self, b):
137137
except ValueError:
138138
pass
139139

140+
# 12. the oversized path calls back into Python once per chunk; a raw whose
141+
# write() reallocates the source object's storage must not leave the write
142+
# loop walking freed memory
143+
ba = bytearray(b"A" * 100)
144+
chunks = []
145+
146+
147+
class MutatingWriter:
148+
def write(self, b):
149+
# reallocates ba's backing storage mid-write
150+
ba[0:1] = b"B" * 200000
151+
chunks.append(1)
152+
return 1
153+
154+
155+
w = _io.BufferedWriter(MutatingWriter(), 8)
156+
assert w.write(ba) == 100
157+
# 100 bytes, one accepted per call, until the 8 byte tail is buffered instead
158+
assert len(chunks) == 92, len(chunks)
159+
140160
print("buffered_writer: ok")

‎src/runtime/modules/IOChecks.hpp‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
#pragma once
2+
3+
#include "runtime/Value.hpp"
4+
#include "runtime/ValueError.hpp"
5+
6+
#include <variant>
7+
8+
namespace py {
9+
10+
// bool is_initialized() const - required by check_initialized()
11+
// bool is_detached() const - optional; reported separately when the object was
12+
// initialized once and then detached
13+
template<typename Derived> class IOChecks
14+
{
15+
const Derived &derived() const { return static_cast<const Derived &>(*this); }
16+
17+
public:
18+
PyResult<std::monostate> check_initialized() const
19+
{
20+
if (derived().is_initialized()) { return Ok(std::monostate{}); }
21+
if constexpr (requires(const Derived &d) { d.is_detached(); }) {
22+
if (derived().is_detached()) {
23+
return Err(value_error("raw stream has been detached"));
24+
}
25+
}
26+
return Err(value_error("I/O operation on uninitialized object"));
27+
}
28+
29+
// the guard pair an accessor wants: usable means constructed *and* still open.
30+
// Spelling it once keeps the two checks from being chained in the wrong order.
31+
PyResult<std::monostate> check_usable() const
32+
{
33+
return check_initialized().and_then([this](auto) { return derived().check_closed(); });
34+
}
35+
};
36+
37+
}// namespace py

‎src/runtime/modules/IOModule.cpp‎

Lines changed: 49 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#include "IOChecks.hpp"
12
#include "Modules.hpp"
23
#include "runtime/MemoryError.hpp"
34
#include "runtime/NotImplementedError.hpp"
@@ -799,7 +800,7 @@ class BufferedIOBase : public IOBase
799800

800801
template<typename T>
801802
// requires(std::is_base_of_v<PyObject, T>)
802-
struct Buffered
803+
struct Buffered : IOChecks<Buffered<T>>
803804
{
804805
PyObject *raw{ nullptr };
805806
bool ok{ false };
@@ -814,18 +815,9 @@ struct Buffered
814815

815816
int64_t readahead() const { return valid_readbuffer() ? buffer->in_avail() : 0; }
816817

817-
PyResult<std::monostate> check_initialized() const
818-
{
819-
if (!ok || !raw) {
820-
if (detached) {
821-
return Err(value_error("raw stream has been detached"));
822-
} else {
823-
return Err(value_error("I/O operation on uninitialized object"));
824-
}
825-
}
818+
bool is_initialized() const { return ok && raw; }
826819

827-
return Ok(std::monostate{});
828-
}
820+
bool is_detached() const { return detached; }
829821

830822
PyResult<std::monostate> check_closed(std::string_view err_msg) const
831823
{
@@ -851,7 +843,7 @@ struct Buffered
851843

852844
PyResult<PyObject *> simple_flush() const
853845
{
854-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
846+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
855847

856848
return raw->get_method(PyString::create("flush").unwrap()).and_then([](PyObject *flush) {
857849
return flush->call(PyTuple::create().unwrap(), PyDict::create().unwrap());
@@ -862,7 +854,7 @@ struct Buffered
862854

863855
PyResult<bool> closed() const
864856
{
865-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
857+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
866858
return raw->get_attribute(PyString::create("closed").unwrap())
867859
.and_then([](PyObject *closed) {
868860
return truthy(closed, VirtualMachine::the().interpreter());
@@ -871,7 +863,7 @@ struct Buffered
871863

872864
PyResult<PyObject *> close()
873865
{
874-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
866+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
875867

876868
// FIXME add lock
877869
auto r = closed();
@@ -899,7 +891,7 @@ struct Buffered
899891

900892
PyResult<PyObject *> seekable() const
901893
{
902-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
894+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
903895
return static_cast<const T *>(this)
904896
->raw->get_method(PyString::create("seekable").unwrap())
905897
.and_then([](PyObject *seekable) -> PyResult<PyObject *> {
@@ -909,7 +901,7 @@ struct Buffered
909901

910902
PyResult<PyObject *> writable() const
911903
{
912-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
904+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
913905
return static_cast<const T *>(this)
914906
->raw->get_method(PyString::create("writable").unwrap())
915907
.and_then([](PyObject *writable) -> PyResult<PyObject *> {
@@ -919,7 +911,7 @@ struct Buffered
919911

920912
PyResult<PyObject *> readable() const
921913
{
922-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
914+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
923915
return static_cast<const T *>(this)
924916
->get_method(PyString::create("readable").unwrap())
925917
.and_then([](PyObject *readable) -> PyResult<PyObject *> {
@@ -929,7 +921,7 @@ struct Buffered
929921

930922
PyResult<PyObject *> fileno() const
931923
{
932-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
924+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
933925
return static_cast<const T *>(this)
934926
->raw->get_method(PyString::create("fileno").unwrap())
935927
.and_then([](PyObject *fileno) -> PyResult<PyObject *> {
@@ -939,7 +931,7 @@ struct Buffered
939931

940932
PyResult<PyObject *> isatty() const
941933
{
942-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
934+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
943935
return static_cast<const T *>(this)
944936
->raw->get_method(PyString::create("isatty").unwrap())
945937
.and_then([](PyObject *isatty) -> PyResult<PyObject *> {
@@ -974,7 +966,7 @@ struct Buffered
974966

975967
PyResult<PyObject *> read(int64_t n)
976968
{
977-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
969+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
978970

979971
if (n < -1) { return Err(value_error("read length must be non-negative or -1")); }
980972

@@ -991,7 +983,7 @@ struct Buffered
991983

992984
PyResult<PyObject *> read1(int64_t n)
993985
{
994-
if (auto err = check_initialized(); err.is_err()) return Err(err.unwrap_err());
986+
if (auto err = this->check_initialized(); err.is_err()) return Err(err.unwrap_err());
995987

996988
if (n < 0) {
997989
// TODO: determine actual buffer size
@@ -1498,14 +1490,20 @@ class BufferedWriter
14981490
return PyInteger::create(len);
14991491
}
15001492

1493+
// everything below re-enters Python (flush() and raw_write() both end up calling
1494+
// raw.write()), which can run arbitrary code that reallocates the source object's
1495+
// storage - e.g. a raw whose write() resizes the very bytearray being written. Take a
1496+
// snapshot now so `start` cannot dangle mid-write.
1497+
const std::vector<char> payload{ start, start + len };
1498+
15011499
// the data does not fit: drain our buffer to raw first
15021500
if (auto r = flush(); r.is_err()) { return r; }
15031501

15041502
// write oversized data directly to raw, bypassing the buffer
15051503
size_t written = 0;
15061504
size_t remaining = len;
15071505
while (remaining > m_buffer_size) {
1508-
auto r = raw_write(start + written, remaining);
1506+
auto r = raw_write(payload.data() + written, remaining);
15091507
if (r.is_err()) { return r; }
15101508
if (r.unwrap() == py_none()) {
15111509
// TODO: raise BlockingIOError
@@ -1535,7 +1533,7 @@ class BufferedWriter
15351533

15361534
// buffer the tail, which is now guaranteed to fit
15371535
if (remaining > 0) {
1538-
const auto count = this->buffer->sputn(start + written, remaining);
1536+
const auto count = this->buffer->sputn(payload.data() + written, remaining);
15391537
ASSERT(static_cast<size_t>(count) == remaining);
15401538
m_buffered_bytes += remaining;
15411539
written += remaining;
@@ -2637,8 +2635,12 @@ class IncrementalNewlineDecoder : public PyBaseObject
26372635
}
26382636
};
26392637

2640-
class StringIO : public TextIOBase
2638+
class StringIO
2639+
: public TextIOBase
2640+
, public IOChecks<StringIO>
26412641
{
2642+
friend class IOChecks<StringIO>;
2643+
26422644
friend ::Heap;
26432645

26442646
std::stringstream m_stringstream;
@@ -2721,23 +2723,17 @@ class StringIO : public TextIOBase
27212723

27222724
PyResult<PyObject *> readable() const
27232725
{
2724-
return check_initialized()
2725-
.and_then([this](auto) { return check_closed(); })
2726-
.and_then([](auto) { return Ok(py_true()); });
2726+
return check_usable().and_then([](auto) { return Ok(py_true()); });
27272727
}
27282728

27292729
PyResult<PyObject *> writable() const
27302730
{
2731-
return check_initialized()
2732-
.and_then([this](auto) { return check_closed(); })
2733-
.and_then([](auto) { return Ok(py_true()); });
2731+
return check_usable().and_then([](auto) { return Ok(py_true()); });
27342732
}
27352733

27362734
PyResult<PyObject *> seekable() const
27372735
{
2738-
return check_initialized()
2739-
.and_then([this](auto) { return check_closed(); })
2740-
.and_then([](auto) { return Ok(py_true()); });
2736+
return check_usable().and_then([](auto) { return Ok(py_true()); });
27412737
}
27422738

27432739
PyResult<PyObject *> closed() const
@@ -2748,16 +2744,12 @@ class StringIO : public TextIOBase
27482744

27492745
PyResult<PyObject *> line_buffering() const
27502746
{
2751-
return check_initialized()
2752-
.and_then([this](auto) { return check_closed(); })
2753-
.and_then([](auto) { return Ok(py_false()); });
2747+
return check_usable().and_then([](auto) { return Ok(py_false()); });
27542748
}
27552749

27562750
PyResult<PyObject *> newlines() const
27572751
{
2758-
return check_initialized()
2759-
.and_then([this](auto) { return check_closed(); })
2760-
.and_then([](auto) { return Ok(py_none()); });
2752+
return check_usable().and_then([](auto) { return Ok(py_none()); });
27612753
}
27622754

27632755
PyResult<PyObject *> close()
@@ -2769,19 +2761,15 @@ class StringIO : public TextIOBase
27692761

27702762
PyResult<PyObject *> tell()
27712763
{
2772-
return check_initialized()
2773-
.and_then([this](auto) { return check_closed(); })
2774-
.and_then([this](auto) { return PyInteger::create(m_stringstream.tellg()); });
2764+
return check_usable().and_then(
2765+
[this](auto) { return PyInteger::create(m_stringstream.tellg()); });
27752766
}
27762767

27772768
PyResult<PyObject *> read(PyTuple *args, PyDict *kwargs)
27782769
{
27792770
ASSERT(!kwargs || kwargs->size() == 0);
27802771

2781-
if (auto result = check_initialized().or_else([this](auto) { return check_closed(); });
2782-
result.is_err()) {
2783-
return result;
2784-
}
2772+
if (auto result = check_usable(); result.is_err()) { return Err(result.unwrap_err()); }
27852773

27862774
auto size_ = [args, this]() -> PyResult<size_t> {
27872775
const auto initial_pos = m_stringstream.tellg();
@@ -2817,10 +2805,7 @@ class StringIO : public TextIOBase
28172805
{
28182806
ASSERT(!kwargs || kwargs->size() == 0);
28192807

2820-
if (auto result = check_initialized().or_else([this](auto) { return check_closed(); });
2821-
result.is_err()) {
2822-
return result;
2823-
}
2808+
if (auto result = check_usable(); result.is_err()) { return Err(result.unwrap_err()); }
28242809

28252810
auto limit_ = [args, this]() -> PyResult<size_t> {
28262811
const auto initial_pos = m_stringstream.tellg();
@@ -2864,7 +2849,7 @@ class StringIO : public TextIOBase
28642849

28652850
PyResult<PyObject *> write(PyTuple *args, PyDict *kwargs)
28662851
{
2867-
if (auto result = check_initialized(); result.is_err()) { return result; }
2852+
if (auto result = check_initialized(); result.is_err()) { return Err(result.unwrap_err()); }
28682853

28692854
auto parse_result = PyArgsParser<PyString *>::unpack_tuple(args,
28702855
kwargs,
@@ -2876,7 +2861,7 @@ class StringIO : public TextIOBase
28762861

28772862
auto [obj] = parse_result.unwrap();
28782863

2879-
if (auto result = check_closed(); result.is_err()) { return result; }
2864+
if (auto result = check_closed(); result.is_err()) { return Err(result.unwrap_err()); }
28802865

28812866
const auto size = obj->size();
28822867
m_stringstream << obj->value();
@@ -2889,7 +2874,7 @@ class StringIO : public TextIOBase
28892874
{
28902875
ASSERT(!kwargs || kwargs->size() == 0);
28912876

2892-
if (auto result = check_initialized(); result.is_err()) { return result; }
2877+
if (auto result = check_initialized(); result.is_err()) { return Err(result.unwrap_err()); }
28932878

28942879
auto parse_result = PyArgsParser<PyInteger *, PyInteger *>::unpack_tuple(args,
28952880
kwargs,
@@ -2970,23 +2955,22 @@ class StringIO : public TextIOBase
29702955
return Ok(0);
29712956
}
29722957

2973-
PyResult<PyObject *> check_initialized() const
2974-
{
2975-
if (!m_ok) { return Err(value_error("I/O operation on uninitialized object")); }
2976-
return Ok(py_true());
2977-
}
2958+
bool is_initialized() const { return m_ok; }
29782959

2979-
PyResult<PyObject *> check_closed() const
2960+
PyResult<std::monostate> check_closed() const
29802961
{
29812962
if (m_closed) { return Err(value_error("I/O operation on closed file")); }
2982-
return Ok(py_true());
2963+
return Ok(std::monostate{});
29832964
}
29842965
};
29852966

29862967

2987-
class TextIOWrapper : public TextIOBase
2968+
class TextIOWrapper
2969+
: public TextIOBase
2970+
, public IOChecks<TextIOWrapper>
29882971
{
29892972
friend ::Heap;
2973+
friend class IOChecks<TextIOWrapper>;
29902974

29912975
PyObject *m_buffer{ nullptr };
29922976
std::string m_errors;
@@ -3327,6 +3311,7 @@ class TextIOWrapper : public TextIOBase
33273311

33283312
PyResult<PyObject *> isatty()
33293313
{
3314+
if (auto err = check_initialized(); err.is_err()) { return Err(err.unwrap_err()); }
33303315
return m_buffer->get_method(PyString::create("isatty").unwrap()).and_then([](auto *isatty) {
33313316
return isatty->call(nullptr, nullptr);
33323317
});
@@ -3376,6 +3361,8 @@ class TextIOWrapper : public TextIOBase
33763361

33773362
return Ok(1);
33783363
}
3364+
3365+
bool is_initialized() const { return m_buffer != nullptr; }
33793366
};
33803367

33813368
PyResult<PyObject *> open(PyObject *file, const std::string &mode)

0 commit comments

Comments
 (0)