1aec3c8f4Simarom"""0MQ Context class."""
2aec3c8f4Simarom# coding: utf-8
3aec3c8f4Simarom
4aec3c8f4Simarom# Copyright (c) PyZMQ Developers.
5aec3c8f4Simarom# Distributed under the terms of the Lesser GNU Public License (LGPL).
6aec3c8f4Simarom
7aec3c8f4Simaromfrom libc.stdlib cimport free, malloc, realloc
8aec3c8f4Simarom
9aec3c8f4Simaromfrom libzmq cimport *
10aec3c8f4Simarom
11aec3c8f4Simaromcdef extern from "getpid_compat.h":
12aec3c8f4Simarom    int getpid()
13aec3c8f4Simarom
14aec3c8f4Simaromfrom zmq.error import ZMQError
15aec3c8f4Simaromfrom zmq.backend.cython.checkrc cimport _check_rc
16aec3c8f4Simarom
17aec3c8f4Simarom
18aec3c8f4Simarom_instance = None
19aec3c8f4Simarom
20aec3c8f4Simaromcdef class Context:
21aec3c8f4Simarom    """Context(io_threads=1)
22aec3c8f4Simarom
23aec3c8f4Simarom    Manage the lifecycle of a 0MQ context.
24aec3c8f4Simarom
25aec3c8f4Simarom    Parameters
26aec3c8f4Simarom    ----------
27aec3c8f4Simarom    io_threads : int
28aec3c8f4Simarom        The number of IO threads.
29aec3c8f4Simarom    """
30aec3c8f4Simarom    
31aec3c8f4Simarom    # no-op for the signature
32aec3c8f4Simarom    def __init__(self, io_threads=1, shadow=0):
33aec3c8f4Simarom        pass
34aec3c8f4Simarom    
35aec3c8f4Simarom    def __cinit__(self, int io_threads=1, size_t shadow=0, **kwargs):
36aec3c8f4Simarom        self.handle = NULL
37aec3c8f4Simarom        self._sockets = NULL
38aec3c8f4Simarom        if shadow:
39aec3c8f4Simarom            self.handle = <void *>shadow
40aec3c8f4Simarom            self._shadow = True
41aec3c8f4Simarom        else:
42aec3c8f4Simarom            self._shadow = False
43aec3c8f4Simarom            if ZMQ_VERSION_MAJOR >= 3:
44aec3c8f4Simarom                self.handle = zmq_ctx_new()
45aec3c8f4Simarom            else:
46aec3c8f4Simarom                self.handle = zmq_init(io_threads)
47aec3c8f4Simarom        
48aec3c8f4Simarom        if self.handle == NULL:
49aec3c8f4Simarom            raise ZMQError()
50aec3c8f4Simarom        
51aec3c8f4Simarom        cdef int rc = 0
52aec3c8f4Simarom        if ZMQ_VERSION_MAJOR >= 3 and not self._shadow:
53aec3c8f4Simarom            rc = zmq_ctx_set(self.handle, ZMQ_IO_THREADS, io_threads)
54aec3c8f4Simarom            _check_rc(rc)
55aec3c8f4Simarom        
56aec3c8f4Simarom        self.closed = False
57aec3c8f4Simarom        self._n_sockets = 0
58aec3c8f4Simarom        self._max_sockets = 32
59aec3c8f4Simarom        
60aec3c8f4Simarom        self._sockets = <void **>malloc(self._max_sockets*sizeof(void *))
61aec3c8f4Simarom        if self._sockets == NULL:
62aec3c8f4Simarom            raise MemoryError("Could not allocate _sockets array")
63aec3c8f4Simarom        
64aec3c8f4Simarom        self._pid = getpid()
65aec3c8f4Simarom    
66aec3c8f4Simarom    def __dealloc__(self):
67aec3c8f4Simarom        """don't touch members in dealloc, just cleanup allocations"""
68aec3c8f4Simarom        cdef int rc
69aec3c8f4Simarom        if self._sockets != NULL:
70aec3c8f4Simarom            free(self._sockets)
71aec3c8f4Simarom            self._sockets = NULL
72aec3c8f4Simarom            self._n_sockets = 0
73aec3c8f4Simarom
74aec3c8f4Simarom        # we can't call object methods in dealloc as it
75aec3c8f4Simarom        # might already be partially deleted
76aec3c8f4Simarom        if not self._shadow:
77aec3c8f4Simarom            self._term()
78aec3c8f4Simarom    
79aec3c8f4Simarom    cdef inline void _add_socket(self, void* handle):
80aec3c8f4Simarom        """Add a socket handle to be closed when Context terminates.
81aec3c8f4Simarom        
82aec3c8f4Simarom        This is to be called in the Socket constructor.
83aec3c8f4Simarom        """
84aec3c8f4Simarom        if self._n_sockets >= self._max_sockets:
85aec3c8f4Simarom            self._max_sockets *= 2
86aec3c8f4Simarom            self._sockets = <void **>realloc(self._sockets, self._max_sockets*sizeof(void *))
87aec3c8f4Simarom            if self._sockets == NULL:
88aec3c8f4Simarom                raise MemoryError("Could not reallocate _sockets array")
89aec3c8f4Simarom        
90aec3c8f4Simarom        self._sockets[self._n_sockets] = handle
91aec3c8f4Simarom        self._n_sockets += 1
92aec3c8f4Simarom
93aec3c8f4Simarom    cdef inline void _remove_socket(self, void* handle):
94aec3c8f4Simarom        """Remove a socket from the collected handles.
95aec3c8f4Simarom        
96aec3c8f4Simarom        This should be called by Socket.close, to prevent trying to
97aec3c8f4Simarom        close a socket a second time.
98aec3c8f4Simarom        """
99aec3c8f4Simarom        cdef bint found = False
100aec3c8f4Simarom        
101aec3c8f4Simarom        for idx in range(self._n_sockets):
102aec3c8f4Simarom            if self._sockets[idx] == handle:
103aec3c8f4Simarom                found=True
104aec3c8f4Simarom                break
105aec3c8f4Simarom        
106aec3c8f4Simarom        if found:
107aec3c8f4Simarom            self._n_sockets -= 1
108aec3c8f4Simarom            if self._n_sockets:
109aec3c8f4Simarom                # move last handle to closed socket's index
110aec3c8f4Simarom                self._sockets[idx] = self._sockets[self._n_sockets]
111aec3c8f4Simarom    
112aec3c8f4Simarom    
113aec3c8f4Simarom    @property
114aec3c8f4Simarom    def underlying(self):
115aec3c8f4Simarom        """The address of the underlying libzmq context"""
116aec3c8f4Simarom        return <size_t> self.handle
117aec3c8f4Simarom    
118aec3c8f4Simarom    # backward-compat, though nobody is using it
119aec3c8f4Simarom    _handle = underlying
120aec3c8f4Simarom    
121aec3c8f4Simarom    cdef inline int _term(self):
122aec3c8f4Simarom        cdef int rc=0
123aec3c8f4Simarom        if self.handle != NULL and not self.closed and getpid() == self._pid:
124aec3c8f4Simarom            with nogil:
125aec3c8f4Simarom                rc = zmq_ctx_destroy(self.handle)
126aec3c8f4Simarom        self.handle = NULL
127aec3c8f4Simarom        return rc
128aec3c8f4Simarom    
129aec3c8f4Simarom    def term(self):
130aec3c8f4Simarom        """ctx.term()
131aec3c8f4Simarom
132aec3c8f4Simarom        Close or terminate the context.
133aec3c8f4Simarom        
134aec3c8f4Simarom        This can be called to close the context by hand. If this is not called,
135aec3c8f4Simarom        the context will automatically be closed when it is garbage collected.
136aec3c8f4Simarom        """
137aec3c8f4Simarom        cdef int rc
138aec3c8f4Simarom        rc = self._term()
139aec3c8f4Simarom        self.closed = True
140aec3c8f4Simarom    
141aec3c8f4Simarom    def set(self, int option, optval):
142aec3c8f4Simarom        """ctx.set(option, optval)
143aec3c8f4Simarom
144aec3c8f4Simarom        Set a context option.
145aec3c8f4Simarom
146aec3c8f4Simarom        See the 0MQ API documentation for zmq_ctx_set
147aec3c8f4Simarom        for details on specific options.
148aec3c8f4Simarom        
149aec3c8f4Simarom        .. versionadded:: libzmq-3.2
150aec3c8f4Simarom        .. versionadded:: 13.0
151aec3c8f4Simarom
152aec3c8f4Simarom        Parameters
153aec3c8f4Simarom        ----------
154aec3c8f4Simarom        option : int
155aec3c8f4Simarom            The option to set.  Available values will depend on your
156aec3c8f4Simarom            version of libzmq.  Examples include::
157aec3c8f4Simarom            
158aec3c8f4Simarom                zmq.IO_THREADS, zmq.MAX_SOCKETS
159aec3c8f4Simarom        
160aec3c8f4Simarom        optval : int
161aec3c8f4Simarom            The value of the option to set.
162aec3c8f4Simarom        """
163aec3c8f4Simarom        cdef int optval_int_c
164aec3c8f4Simarom        cdef int rc
165aec3c8f4Simarom        cdef char* optval_c
166aec3c8f4Simarom
167aec3c8f4Simarom        if self.closed:
168aec3c8f4Simarom            raise RuntimeError("Context has been destroyed")
169aec3c8f4Simarom        
170aec3c8f4Simarom        if not isinstance(optval, int):
171aec3c8f4Simarom            raise TypeError('expected int, got: %r' % optval)
172aec3c8f4Simarom        optval_int_c = optval
173aec3c8f4Simarom        rc = zmq_ctx_set(self.handle, option, optval_int_c)
174aec3c8f4Simarom        _check_rc(rc)
175aec3c8f4Simarom
176aec3c8f4Simarom    def get(self, int option):
177aec3c8f4Simarom        """ctx.get(option)
178aec3c8f4Simarom
179aec3c8f4Simarom        Get the value of a context option.
180aec3c8f4Simarom
181aec3c8f4Simarom        See the 0MQ API documentation for zmq_ctx_get
182aec3c8f4Simarom        for details on specific options.
183aec3c8f4Simarom        
184aec3c8f4Simarom        .. versionadded:: libzmq-3.2
185aec3c8f4Simarom        .. versionadded:: 13.0
186aec3c8f4Simarom
187aec3c8f4Simarom        Parameters
188aec3c8f4Simarom        ----------
189aec3c8f4Simarom        option : int
190aec3c8f4Simarom            The option to get.  Available values will depend on your
191aec3c8f4Simarom            version of libzmq.  Examples include::
192aec3c8f4Simarom            
193aec3c8f4Simarom                zmq.IO_THREADS, zmq.MAX_SOCKETS
194aec3c8f4Simarom            
195aec3c8f4Simarom        Returns
196aec3c8f4Simarom        -------
197aec3c8f4Simarom        optval : int
198aec3c8f4Simarom            The value of the option as an integer.
199aec3c8f4Simarom        """
200aec3c8f4Simarom        cdef int optval_int_c
201aec3c8f4Simarom        cdef size_t sz
202aec3c8f4Simarom        cdef int rc
203aec3c8f4Simarom
204aec3c8f4Simarom        if self.closed:
205aec3c8f4Simarom            raise RuntimeError("Context has been destroyed")
206aec3c8f4Simarom
207aec3c8f4Simarom        rc = zmq_ctx_get(self.handle, option)
208aec3c8f4Simarom        _check_rc(rc)
209aec3c8f4Simarom
210aec3c8f4Simarom        return rc
211aec3c8f4Simarom
212aec3c8f4Simarom    def destroy(self, linger=None):
213aec3c8f4Simarom        """ctx.destroy(linger=None)
214aec3c8f4Simarom        
215aec3c8f4Simarom        Close all sockets associated with this context, and then terminate
216aec3c8f4Simarom        the context. If linger is specified,
217aec3c8f4Simarom        the LINGER sockopt of the sockets will be set prior to closing.
218aec3c8f4Simarom        
219aec3c8f4Simarom        .. warning::
220aec3c8f4Simarom        
221aec3c8f4Simarom            destroy involves calling ``zmq_close()``, which is **NOT** threadsafe.
222aec3c8f4Simarom            If there are active sockets in other threads, this must not be called.
223aec3c8f4Simarom        """
224aec3c8f4Simarom        
225aec3c8f4Simarom        cdef int linger_c
226aec3c8f4Simarom        cdef bint setlinger=False
227aec3c8f4Simarom        
228aec3c8f4Simarom        if linger is not None:
229aec3c8f4Simarom            linger_c = linger
230aec3c8f4Simarom            setlinger=True
231aec3c8f4Simarom
232aec3c8f4Simarom        if self.handle != NULL and not self.closed and self._n_sockets:
233aec3c8f4Simarom            while self._n_sockets:
234aec3c8f4Simarom                if setlinger:
235aec3c8f4Simarom                    zmq_setsockopt(self._sockets[0], ZMQ_LINGER, &linger_c, sizeof(int))
236aec3c8f4Simarom                rc = zmq_close(self._sockets[0])
237aec3c8f4Simarom                if rc < 0 and zmq_errno() != ZMQ_ENOTSOCK:
238aec3c8f4Simarom                    raise ZMQError()
239aec3c8f4Simarom                self._n_sockets -= 1
240aec3c8f4Simarom                self._sockets[0] = self._sockets[self._n_sockets]
241aec3c8f4Simarom        self.term()
242aec3c8f4Simarom    
243aec3c8f4Simarom__all__ = ['Context']
244