|
17 | 17 | #include"miscadmin.h" |
18 | 18 | #include"pgstat.h" |
19 | 19 | #include"storage/ipc.h" |
| 20 | +#include"storage/procarray.h" |
20 | 21 | #include"storage/procsignal.h" |
21 | 22 | #include"storage/shm_toc.h" |
22 | 23 | #include"utils/guc.h" |
| 24 | +#include"utils/timestamp.h" |
23 | 25 |
|
24 | 26 | #ifdefPG_MODULE_MAGIC |
25 | 27 | PG_MODULE_MAGIC; |
@@ -60,6 +62,7 @@ static void qs_postExecProcNode(PlanState *planstate, TupleTableSlot *result); |
60 | 62 | /* Global variables */ |
61 | 63 | List*QueryDescStack=NIL; |
62 | 64 | staticProcSignalReasonQueryStatePollReason; |
| 65 | +staticProcSignalReasonRolePollReason; |
63 | 66 | staticboolmodule_initialized= false; |
64 | 67 | staticconstchar*be_state_str[]= {/* BackendState -> string repr */ |
65 | 68 | "undefined",/* STATE_UNDEFINED */ |
@@ -90,6 +93,9 @@ typedef struct |
90 | 93 | pid_ttraceable; |
91 | 94 | }trace_request; |
92 | 95 |
|
| 96 | +staticvoidSendCurrentRoleOid(void); |
| 97 | +OidGetRemoteBackendUser(pid_tpid,int*error_code); |
| 98 | + |
93 | 99 | /* Shared memory variables */ |
94 | 100 | shm_toc*toc=NULL; |
95 | 101 | user_data*caller=NULL; |
@@ -181,7 +187,8 @@ _PG_init(void) |
181 | 187 |
|
182 | 188 | /* Register interrupt on custom signal of polling query state */ |
183 | 189 | QueryStatePollReason=RegisterCustomProcSignalHandler(SendQueryState); |
184 | | -if (QueryStatePollReason==INVALID_PROCSIGNAL) |
| 190 | +RolePollReason=RegisterCustomProcSignalHandler(SendCurrentRoleOid); |
| 191 | +if (QueryStatePollReason==INVALID_PROCSIGNAL||RolePollReason==INVALID_PROCSIGNAL) |
185 | 192 | { |
186 | 193 | ereport(WARNING, (errcode(ERRCODE_INSUFFICIENT_RESOURCES), |
187 | 194 | errmsg("pg_query_state isn't loaded: insufficient custom ProcSignal slots"))); |
@@ -608,6 +615,9 @@ pg_query_state(PG_FUNCTION_ARGS) |
608 | 615 | init_lock_tag(&tag,PG_QUERY_STATE_KEY); |
609 | 616 | LockAcquire(&tag,ExclusiveLock, false, false); |
610 | 617 |
|
| 618 | +interror_code; |
| 619 | +Oiduser_id=GetRemoteBackendUser(pid,&error_code); |
| 620 | + |
611 | 621 | /* fill in caller's user data */ |
612 | 622 | caller->user_id=GetUserId(); |
613 | 623 | caller->superuser=superuser(); |
@@ -811,3 +821,87 @@ executor_continue(PG_FUNCTION_ARGS) |
811 | 821 |
|
812 | 822 | PG_RETURN_VOID(); |
813 | 823 | } |
| 824 | + |
| 825 | +staticshm_mq_result |
| 826 | +shm_mq_receive_with_timeout(shm_mq_handle*mqh,Size*nbytesp,void**datap,longtimeout) |
| 827 | +{ |
| 828 | + |
| 829 | +#ifdefHAVE_INT64_TIMESTAMP |
| 830 | +#defineGetNowLong()((long) GetCurrentTimestamp() / 1000) |
| 831 | +#else |
| 832 | +#defineGetNowLong()1000 * GetCurrentTimestamp() |
| 833 | +#endif |
| 834 | + |
| 835 | +longendtime=GetNowLong()+timeout; |
| 836 | + |
| 837 | +for (;;) |
| 838 | +{ |
| 839 | +intrc; |
| 840 | +longdelay; |
| 841 | +shm_mq_resultmq_receive_result=shm_mq_receive(mqh,nbytesp,datap, true); |
| 842 | + |
| 843 | +if (mq_receive_result!=SHM_MQ_WOULD_BLOCK) |
| 844 | +returnmq_receive_result; |
| 845 | + |
| 846 | +delay=endtime-GetNowLong(); |
| 847 | +rc=WaitLatch(MyLatch,WL_LATCH_SET |WL_TIMEOUT,delay); |
| 848 | +CHECK_FOR_INTERRUPTS(); |
| 849 | +ResetLatch(MyLatch); |
| 850 | + |
| 851 | +if (rc&WL_TIMEOUT) |
| 852 | +returnSHM_MQ_WOULD_BLOCK; |
| 853 | +} |
| 854 | +} |
| 855 | + |
| 856 | +staticvoid |
| 857 | +SendCurrentRoleOid(void) |
| 858 | +{ |
| 859 | +shm_mq_handle*mqh=shm_mq_attach(mq,NULL,NULL); |
| 860 | +Oidrole_oid=GetUserId(); |
| 861 | + |
| 862 | +shm_mq_send(mqh,sizeof(Oid),&role_oid, false); |
| 863 | +} |
| 864 | + |
| 865 | +#defineNOT_BACKEND_PROCESS1 |
| 866 | +#defineCOULD_NOT_SEND_SIGNAL 2 |
| 867 | +#defineINVALID_MQ_READ3 |
| 868 | + |
| 869 | +Oid |
| 870 | +GetRemoteBackendUser(pid_tpid,int*error_code) |
| 871 | +{ |
| 872 | +PGPROC*proc=BackendPidGetProc(pid); |
| 873 | +intsig_result; |
| 874 | +shm_mq_handle*mqh; |
| 875 | +shm_mq_resultmq_receive_result; |
| 876 | +Oid*result; |
| 877 | +Sizeres_len; |
| 878 | + |
| 879 | +if (proc==NULL||proc->backendId==InvalidBackendId) |
| 880 | +{ |
| 881 | +*error_code=NOT_BACKEND_PROCESS; |
| 882 | +returnInvalidOid; |
| 883 | +} |
| 884 | + |
| 885 | +mq=shm_mq_create(mq,QUEUE_SIZE); |
| 886 | +shm_mq_set_sender(mq,proc); |
| 887 | +shm_mq_set_receiver(mq,MyProc); |
| 888 | + |
| 889 | +sig_result=SendProcSignal(pid,RolePollReason,proc->backendId); |
| 890 | +if (sig_result==-1) |
| 891 | +{ |
| 892 | +*error_code=COULD_NOT_SEND_SIGNAL; |
| 893 | +returnInvalidOid; |
| 894 | +} |
| 895 | + |
| 896 | +mqh=shm_mq_attach(mq,NULL,NULL); |
| 897 | +mq_receive_result=shm_mq_receive_with_timeout(mqh,&res_len, (void**)&result,1000); |
| 898 | +if (mq_receive_result!=SHM_MQ_SUCCESS) |
| 899 | +{ |
| 900 | +*error_code=INVALID_MQ_READ; |
| 901 | +returnInvalidOid; |
| 902 | +} |
| 903 | + |
| 904 | +shm_mq_detach(mq); |
| 905 | + |
| 906 | +return*result; |
| 907 | +} |
