Fixing memory leaks in popular Python libraries | Hacker News Reader

Fixing memory leaks in popular Python libraries

paulsprogrammingnotes.com

128 pointspostpawl35 comments

FWIW, Python will close a socket automatically when the reference to the socket is GC'd. (Of course, you _should_ always close a socket explicitly and not rely on GC to handle it for you.)

So the larger issue here was that the original code was written such that when shutdown() raised an exception, the reference to the socket object wasn't released—i.e., the line of code that set `self.socket = None` was not reached.

I had to follow the link to the PR to see what was going on, because this wasn't clear to me from the minimal example in the blog post. The fix to redis-py wasn't 100% necessary. Here's the redis-py code before the PR:

https://github.com/pawl/redis-py/blob/82bad1686177c4c543818a...

   def disconnect(self):
        "Disconnects from the Redis server"
        self._parser.on_disconnect()
        if self._sock is None:
            return
        try:
            if os.getpid() == self.pid:
                self._sock.shutdown(socket.SHUT_RDWR)
            self._sock.close()
        except OSError:
            pass
        self._sock = None

Notice that self._sock is set to None even if shutdown() raises an exception. This will eventually cause the socket object to be GC'd, and thus closed. No memory leak here. But it's still good to call close explicitly, so I'm not disagree'ing with the PR, just that there wasn't a memory leak:

https://github.com/pawl/redis-py/commit/d7f4ac33217a95541ddf...

Conversely, py-amqp looked like this:

    def close(self):
        if self.sock is not None:
            self._shutdown_transport()
            # Call shutdown first to make sure that pending messages
            # reach the AMQP broker if the program exits after
            # calling this method.
            self.sock.shutdown(socket.SHUT_RDWR)
            self.sock.close()
            self.sock = None
        self.connected = False

Any exception in this method prevents `self.sock = None` from running. That, _combined with the fact that_ Celery never released the py-amqp object is where the leak was.

I tested your theory on a branch of py-amqp: https://github.com/celery/celery/issues/4843#issuecomment-99...

It still leaks if you don't close the socket.

Interesting. I conjecture this is because Python doesn’t gc the socket object immediately. Python is supposed to close file-descriptor like objects when they are gc’d, so I'd really like to better understand what's going on. Thank you for all the work so far. I'll try to figure out why my theory isn't empirically proving true.

Edit: I've just done a simple test with Python 3.9.5 and it definitely closes a socket when the object is GC'd. In one terminal, I did this:

   nc -l 9999

In a second terminal:

    Python 3.9.5 (default, May  3 2021, 19:12:05)
    [Clang 12.0.5 (clang-1205.0.22.9)] on darwin
    Type "help", "copyright", "credits" or "license" for more information.
    >>> import socket
    >>> s = socket.socket()
    >>> s.connect(("127.0.0.1", 9999))
    >>> s = None

As soon as `s = None` is executed, `nc` exits. Here's what the full sequence looks like in tcpdump. First the three-way handshake (SYN, ACK, SYN) when the socket is connected:

    13:09:24.448234 IP 127.0.0.1.62086 > 127.0.0.1.9999: Flags [S], seq 2594347262, win 65535, options [mss 16344,nop,wscale 6,nop,nop,TS val 62632264 ecr 0,sackOK,eol], length 0
    13:09:24.448385 IP 127.0.0.1.9999 > 127.0.0.1.62086: Flags [S.], seq 682643237, ack 2594347263, win 65535, options [mss 16344,nop,wscale 6,nop,nop,TS val 3095724505 ecr 62632264,sackOK,eol], length 0
   13:09:24.448410 IP 127.0.0.1.62086 > 127.0.0.1.9999: Flags [.], ack 1, win 6379, options [nop,nop,TS val 62632264 ecr 3095724505], length 0
   13:09:24.448422 IP 127.0.0.1.9999 > 127.0.0.1.62086: Flags [.], ack 1, win 6379, options [nop,nop,TS val 3095724505 ecr 62632264], length 0

Then the active close when `s = None` executes:

    13:09:28.234099 IP 127.0.0.1.62086 > 127.0.0.1.9999: Flags [F.], seq 1, ack 1, win 6379, options [nop,nop,TS val 62636050 ecr 3095724505], length 0
    13:09:28.234178 IP 127.0.0.1.9999 > 127.0.0.1.62086: Flags [.], ack 2, win 6379, options [nop,nop,TS val 3095728291 ecr 62636050], length 0
    13:09:28.234203 IP 127.0.0.1.9999 > 127.0.0.1.62086: Flags [F.], seq 1, ack 2, win 6379, options [nop,nop,TS val 3095728291 ecr 62636050], length 0
   13:09:28.234240 IP 127.0.0.1.62086 > 127.0.0.1.9999: Flags [.], ack 2, win 6379, options [nop,nop,TS val 62636050 ecr 3095728291], length 0

We see the FINs from both sides along with the final ACKs.

So, I don't know what the heck is causing this apparent leak. Just setting `self.sock = None` should be sufficient, modulo the vagaries of when Python GC occurs on the object.

You just need 1 other reference to the socket in code to prevent it from being garbage collected without socket.close().

For example, this won't automatically close the socket:

    >>> s = socket.socket()
    >>> other_sock = s
    >>> s.connect(("127.0.0.1", 9999))
    >>> s = None

Ah, of course. Possibly this then:

    def _setup_transport(self):
        # Setup to _write() directly to the socket, and
        # do our own buffered reads.
        self._write = self.sock.sendall
        self._read_buffer = EMPTY_BUFFER
        self._quick_recv = self.sock.recv

(And similarly in SSLTransport.)

Then, the OSError exception raised inside Transport.close also prevents Connection.collect from destroying the Transport object.

Again, I'm not disagreeing with the original PR/fix. My point is only that when the socket object is destroyed, the socket is closed.

I'll bet a try/except around this (Connection.collect) would also fix it:

    def collect(self):
        if self._transport:
            self._transport.close()

Edit: we're discussing this in two places. My apologies. For interested readers:

https://github.com/celery/celery/issues/4843#issuecomment-99...

I'm satisfied I understand what's going on now, and that it was a bug in celery/py-aqmp retaining the socket object due to over-broad exception handling. Python closes sockets as expected when they are released.

Yeah it should probably be setting “self._quick_recv = None” and “self._write = None” when it shuts down the transport. This probably would have fixed the leak too.

I wasn’t seeing any leaky behavior after the fix though, so I think calling “socket.close()” allows GC to still clean up those other references.

Right, at that point the entire Transport object is destroyed by collect.

Seems like a context manager is what you really want, right? some other exception could be thrown from some deeper frame while sock is in scope and this wouldn't ensure it gets closed.

Or a "finally" block.

The sock.close() can also throw it’s own OSError, so it needs its own try/except block.

For example, this wouldn’t handle the OSError on close:

    try:
        sock.shutdown()
    except OSError:
        pass
    finally:
        sock.close()

    try:
        try:
            sock.shutdown()
        finally:
            sock.close()
    except OSError:
        pass

Or, given the .close method

    from contextlib import closing

    try:    
        with closing(sock):
           sock.shutdown()
    except OSError:
        pass

`contextlib.ignore` could also be useful here e.g.

    with contextlib.ignore(OSError),\
         closing(sock):
        sock.shutdown()

contextlib.ignore was renamed to contextlib.suppress

There is nothing wrong with the proposed solution though? I think you are over doing it IMHO, creating unnecessary complexion.

I’ve no idea what you’re to trying to express.

My comment does not imply there’s anything wrong with the comment, it just uses a tool from the standard library to reduce nesting and (as far as i’m concerned) make it more explicit and upfront.

And support the original query.

well, the nested try: ... finally: gave you something to think about: why do you have a "sock" object that you need to explicitly close "manually"?

I didn't look at the source, it can be for a number of reasons. But I suspect some long lived sock object shared between functions.

Something that could be wrapped in a context manager.

    with MyCtxMgr() as my_ctx_mgr:              # open the socket, init, etc...
       my_ctx_mgr.do_stuff()                    # internally it has access
       my_ctx_mgr.do_another_stuff()            # to the opened socket
       external_function(my_ctx_mgr.get_sock()) # and can pass it

    # in the __exit__ actually close the socket, dealing with the OSError

calling a barebone "obj.close()" is a code smell IMHO

As far as I know, the only errors that you're supposed to "except: pass" from socket.shutdown are OSErrors. I think you want anything else (if it even throws anything else?) to be fatal.

This is the same way cpython handles it: https://github.com/python/cpython/blob/b56774bd93822e1598bb5...

The context manager idea sounds good, but I think that might require a lot of refactoring.

Way to go! Great idea checking other projects too to see if they could also benefit from the fix.

Shouldn't

   sock.close()

imply

   sock.shutdown()

? Isn't it pointless to call both?

https://beej.us/guide/bgnet/html/#close-and-shutdownget-outt...

You're 100% right... In this case the shutdown is called with RDWR... so the author could simply call close!

Is there any case where a shutdown() right before a close() makes a difference?

Why would they call shutdown and then immediately close?

Just call close immediately.

In fact, there is almost no need for anyone ever to call shutdown.

It may make sense if you want to pass a socket descriptor to another process and you only want them to read from the socket, not write to it, to call a unidirectional shutdown on it before passing it on. But shutdown RDWR and then close is just silly.

The proper fix is not to add more exception handling but to remove the shutdown line altogether.

This SO post describes how shutdown is different than close: https://stackoverflow.com/questions/409783/socket-shutdown-v...

So maybe shutdown helps in case other processes are doing something with the socket? I saw this same shutdown + close pattern in cpython code, but urllib3 doesn’t call shutdown before closing the socket (they have a TODO about it). I’m also curious if calling shutdown is necessary.

The memory optimization at the end is neat

I assume there are generally only a few channel ids being used, since otherwise the code could've still been optimized to use a bitset instead of an array of numbers. That'd only require 8192 bytes, a sixteenth of previous code, but still quite a bit more than the 64 bytes an empty array uses

I understand how the example can leak a resource from the OS standpoint, but I don’t understand how this ends up being a memory leak for the Python process.

It’s not the first time I am hearing about memleaks for a GC-ed language, and I never really understood how it was possible. Would anyone have a super simple example?

Not sure if that classifies as a memory leak,but it can happen with mutable types (lists for example) attached to some global variable or as a default in a class definition / function input.

For example:

  class A
      a = []
      def __init__(self):
          self.a.append(1)

  def foo(a=[]):
      # code that adds to a...

Python uses reference counting. The exception was preventing the code from reaching the point where it sets the reference to the socket to None. This isn't clear in the minimized example in the post, but if you follow the link to the PR[1], the original code looked like this:

    def close(self):
        if self.sock is not None:
            self._shutdown_transport()
            # Call shutdown first to make sure that pending messages
            # reach the AMQP broker if the program exits after
            # calling this method.
            self.sock.shutdown(socket.SHUT_RDWR)
            self.sock.close()
            self.sock = None
        self.connected = False

If either shutdown (or close) raises an exception, self.sock isn't set to None. Thus Python retains the reference. Aside, but setting self.sock to None will eventually close the socket once Python GC's the reference[2]. It's still best to explicitly call close though. From [2]:

> Python takes the automatic shutdown a step further, and says that when a socket is garbage collected, it will automatically do a close if it’s needed. But relying on this is a very bad habit. If your socket just disappears without doing a close, the socket at the other end may hang indefinitely, thinking you’re just being slow. Please close your sockets when you’re done.

Another way you can have leaks with a reference counting GC is through circular references. Python provides WeakRef for such situations, but it's up to the programmer to use it correctly.

1. https://github.com/celery/celery/issues/4843#issuecomment-99...

2. https://docs.python.org/3/howto/sockets.html#disconnecting

Python uses a generational reference counting GC, it should be able to handle circular references.

Sheesh, I've somehow gone two decades of Python programming without realizing that. Thank you for the correction.

More info on CPython's GC:

https://devguide.python.org/garbage_collector/

At first I thought this blog was related to Paul's Math Notes, https://tutorial.math.lamar.edu/, but this is a different Paul that lives in a different part of TX.

Hmmmmm shouldn't the Python runtime reap the socket and tidy up in this case?

Yes, I think it should eventually if you haven't stashed a reference somewhere.

Python has both a reference count and a GC so in simple cases it will often be cleaned up as soon as it goes out of scope but if you create a cycle it may take a while. Relying on the GC to close open sockets is a great way to run out of file descriptors when bursts of work happen.

The way the original try/except was structured meant the close() for the socket would never get called, so py-amqp holds onto the reference indefinitely, preventing the garbage collector from doing it's thing.

Reply on news.ycombinator.com